@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e3ab93c896
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 11, 06:15 UTC

[CPU] Refactor CPU fused MOE (#30531)

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

Li, Jiang committed Dec 18, 2025 at 14:36 UTC e3ab93c89667bcd026582e74bc7c49774724a273
23 files changed +1394 -206
.buildkite/scripts/hardware_ci/run-cpu-test.sh
+1
@@ -50,6 +50,7 @@ function cpu_tests() {
50 docker exec cpu-test-"$NUMA_NODE" bash -c "
51 set -e
52 pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
53 + pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
54 pytest -x -v -s tests/kernels/test_onednn.py"
55
56 # Run basic model test
cmake/cpu_extension.cmake
+2 -2
@@ -330,7 +330,7 @@ if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR (ASIMD_FOUND AND NOT APPLE_SILICON
330 PUBLIC ${oneDNN_BINARY_DIR}/include
331 PRIVATE ${oneDNN_SOURCE_DIR}/src
332 )
333 - target_link_libraries(dnnl_ext dnnl)
333 + target_link_libraries(dnnl_ext dnnl torch)
334 target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS} -fPIC)
335 list(APPEND LIBS dnnl_ext)
336 set(USE_ONEDNN ON)
@@ -358,13 +358,13 @@ set(VLLM_EXT_SRC
358 "csrc/cpu/pos_encoding.cpp"
359 "csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
360 "csrc/cpu/cpu_attn.cpp"
361 - "csrc/cpu/scratchpad_manager.cpp"
361 "csrc/cpu/torch_bindings.cpp")
362
363 if (AVX512_FOUND AND NOT AVX512_DISABLED)
364 set(VLLM_EXT_SRC
365 "csrc/cpu/shm.cpp"
366 "csrc/cpu/cpu_wna16.cpp"
367 + "csrc/cpu/cpu_fused_moe.cpp"
368 ${VLLM_EXT_SRC})
369 if (ENABLE_AVX512BF16 AND ENABLE_AVX512VNNI)
370 set(VLLM_EXT_SRC
csrc/cpu/cpu_arch_macros.h renamed
+5 -5
@@ -1,5 +1,5 @@
1 -#ifndef CPU_ATTN_MACROS_H
2 -#define CPU_ATTN_MACROS_H
1 +#ifndef CPU_ARCH_MACROS_H
2 +#define CPU_ARCH_MACROS_H
3
4 // x86_64
5 #ifdef __x86_64__
@@ -26,7 +26,7 @@
26 _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218)); \
27 const __m512i vec_127 = _mm512_set1_epi32(0x0000007f); \
28 const int n_mantissa_bits = 23; \
29 - auto fast_exp = [&](vec_op::FP32Vec16& vec) __attribute__(( \
29 + auto fast_exp = [&](const vec_op::FP32Vec16& vec) __attribute__(( \
30 always_inline)) { \
31 __m512 values = vec.reg; \
32 auto less_ln_flt_min_mask = \
@@ -98,7 +98,7 @@
98 poly = vbslq_f32(hi_mask, inf, poly); \
99 return vbslq_f32(lo_mask, zero, poly); \
100 }; \
101 - auto fast_exp = [&](vec_op::FP32Vec16& vec) \
101 + auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
102 __attribute__((always_inline)) { \
103 float32x4x4_t result; \
104 result.val[0] = neon_expf(vec.reg.val[0]); \
@@ -110,4 +110,4 @@
110
111 #endif // __aarch64__
112
113 -#endif
\ No newline at end of file
113 +#endif
csrc/cpu/cpu_attn_impl.hpp
+9 -32
@@ -8,10 +8,8 @@
8 #include <sys/sysctl.h>
9 #endif
10
11 -#include "cpu_types.hpp"
12 -#include "scratchpad_manager.h"
13 -#include "cpu_attn_macros.h"
14 -#include "utils.hpp"
11 +#include "cpu/cpu_arch_macros.h"
12 +#include "cpu/utils.hpp"
13
14 namespace cpu_attention {
15 enum class ISA { AMX, VEC, VEC16, NEON };
@@ -378,12 +376,13 @@ class AttentionScheduler {
376
377 static constexpr int32_t MaxQTileIterNum = 128;
378
381 - AttentionScheduler() : available_cache_size_(get_available_l2_size()) {}
379 + AttentionScheduler()
380 + : available_cache_size_(cpu_utils::get_available_l2_size()) {}
381
382 torch::Tensor schedule(const ScheduleInput& input) const {
383 const bool casual = input.casual;
384 const int32_t thread_num = omp_get_max_threads();
386 - const int64_t cache_size = get_available_l2_size();
385 + const int64_t cache_size = cpu_utils::get_available_l2_size();
386 const int32_t max_num_q_per_iter = input.max_num_q_per_iter;
387 const int32_t kv_len_alignment = input.kv_block_alignment;
388 int32_t q_head_per_kv = input.num_heads_q / input.num_heads_kv;
@@ -659,7 +658,7 @@ class AttentionScheduler {
658 metadata_ptr->thread_num +
659 metadata_ptr->reduction_scratchpad_size_per_kv_head *
660 (use_gqa ? input.num_heads_kv : input.num_heads_q);
662 - DNNLScratchPadManager::get_dnnl_scratchpad_manager()->realloc(
661 + cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(
662 scratchpad_size);
663
664 // metadata_ptr->print();
@@ -667,7 +666,7 @@ class AttentionScheduler {
666 // test out of boundary access
667 // {
668 // float* cache_ptr =
670 - // DNNLScratchPadManager::get_dnnl_scratchpad_manager()->get_data<float>();
669 + // cpu_utils::ScratchPadManager::getl_scratchpad_manager()->get_data<float>();
670 // for (int64_t i = 0; i < scratchpad_size / sizeof(float); ++i) {
671 // cache_ptr[i] = std::numeric_limits<float>::quiet_NaN();
672 // }
@@ -749,27 +748,6 @@ class AttentionScheduler {
748 return std::max(rounded_tile_size, round_size);
749 }
750
752 - static int64_t get_available_l2_size() {
753 - static int64_t size = []() {
754 -#if defined(__APPLE__)
755 - // macOS doesn't have _SC_LEVEL2_CACHE_SIZE. Use sysctlbyname.
756 - int64_t l2_cache_size = 0;
757 - size_t len = sizeof(l2_cache_size);
758 - if (sysctlbyname("hw.l2cachesize", &l2_cache_size, &len, NULL, 0) == 0 &&
759 - l2_cache_size > 0) {
760 - return l2_cache_size >> 1; // use 50% of L2 cache
761 - }
762 - // Fallback if sysctlbyname fails
763 - return 128LL * 1024 >> 1; // use 50% of 128KB
764 -#else
765 - long l2_cache_size = sysconf(_SC_LEVEL2_CACHE_SIZE);
766 - TORCH_CHECK_NE(l2_cache_size, -1);
767 - return l2_cache_size >> 1; // use 50% of L2 cache
768 -#endif
769 - }();
770 - return size;
771 - }
772 -
751 private:
752 int64_t available_cache_size_;
753 };
@@ -1402,7 +1380,7 @@ class AttentionMainLoop {
1380
1381 // init buffers
1382 void* scratchpad_ptr =
1405 - DNNLScratchPadManager::get_dnnl_scratchpad_manager()
1383 + cpu_utils::ScratchPadManager::get_scratchpad_manager()
1384 ->get_data<void>();
1385 AttentionScratchPad buffer_manager(thread_id, metadata, scratchpad_ptr);
1386
@@ -1422,8 +1400,7 @@ class AttentionMainLoop {
1400 }
1401 }
1402
1425 - const int64_t available_cache_size =
1426 - AttentionScheduler::get_available_l2_size();
1403 + const int64_t available_cache_size = cpu_utils::get_available_l2_size();
1404 const int32_t default_tile_size =
1405 AttentionScheduler::calcu_default_tile_size(
1406 available_cache_size, head_dim, sizeof(kv_cache_t),
csrc/cpu/cpu_fused_moe.cpp new
+727
@@ -0,0 +1,727 @@
1 +#include "cpu/cpu_types.hpp"
2 +#include "cpu/utils.hpp"
3 +#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
4 +#include "cpu/cpu_arch_macros.h"
5 +
6 +#ifdef CPU_CAPABILITY_AMXBF16
7 + #include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp"
8 + #define AMX_DISPATCH(...) \
9 + case cpu_utils::ISA::AMX: { \
10 + using gemm_t = cpu_micro_gemm::MicroGemm<cpu_utils::ISA::AMX, scalar_t>; \
11 + return __VA_ARGS__(); \
12 + }
13 +#else
14 + #define AMX_DISPATCH(...) case cpu_utils::ISA::AMX:
15 +#endif
16 +
17 +#define CPU_ISA_DISPATCH_IMPL(ISA_TYPE, ...) \
18 + [&] { \
19 + switch (ISA_TYPE) { \
20 + AMX_DISPATCH(__VA_ARGS__) \
21 + case cpu_utils::ISA::VEC: { \
22 + using gemm_t = \
23 + cpu_micro_gemm::MicroGemm<cpu_utils::ISA::VEC, scalar_t>; \
24 + return __VA_ARGS__(); \
25 + } \
26 + default: { \
27 + TORCH_CHECK(false, "Invalid CPU ISA type."); \
28 + } \
29 + } \
30 + }()
31 +
32 +namespace {
33 +enum class FusedMOEAct { SiluAndMul, SwigluOAIAndMul };
34 +
35 +FusedMOEAct get_act_type(const std::string& act) {
36 + if (act == "silu") {
37 + return FusedMOEAct::SiluAndMul;
38 + } else if (act == "swigluoai") {
39 + return FusedMOEAct::SwigluOAIAndMul;
40 + } else {
41 + TORCH_CHECK(false, "Invalid act type: " + act);
42 + }
43 +}
44 +
45 +template <typename scalar_t>
46 +void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
47 + const int32_t m_size, const int32_t n_size,
48 + const int32_t input_stride,
49 + const int32_t output_stride) {
50 + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
51 + // For GPT-OSS interleaved gate-up weights
52 + alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14,
53 + 16, 18, 20, 22, 24, 26, 28, 30};
54 + vec_op::INT32Vec16 index_vec(index);
55 + vec_op::FP32Vec16 gate_up_max_vec(7.0);
56 + vec_op::FP32Vec16 up_min_vec(-7.0);
57 + vec_op::FP32Vec16 alpha_vec(1.702);
58 + vec_op::FP32Vec16 one_vec(1.0);
59 +
60 + DEFINE_FAST_EXP
61 +
62 + for (int32_t m = 0; m < m_size; ++m) {
63 + for (int32_t n = 0; n < n_size; n += 32) {
64 + vec_op::FP32Vec16 gate_vec(input + n, index_vec);
65 + vec_op::FP32Vec16 up_vec(input + n + 1, index_vec);
66 + gate_vec = gate_vec.min(gate_up_max_vec);
67 + up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec);
68 + auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec));
69 + auto glu = gate_vec * sigmoid_vec;
70 + auto gated_output_fp32 = (one_vec + up_vec) * glu;
71 + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
72 + gated_output.save(output + n / 2);
73 + }
74 + input += input_stride;
75 + output += output_stride;
76 + }
77 +}
78 +
79 +template <typename scalar_t>
80 +void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
81 + const int32_t m_size, const int32_t n_size,
82 + const int32_t input_stride, const int32_t output_stride) {
83 + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
84 + const int32_t dim = n_size / 2;
85 + float* __restrict__ gate = input;
86 + float* __restrict__ up = input + dim;
87 + vec_op::FP32Vec16 one_vec(1.0);
88 +
89 + DEFINE_FAST_EXP
90 +
91 + for (int32_t m = 0; m < m_size; ++m) {
92 + for (int32_t n = 0; n < dim; n += 16) {
93 + vec_op::FP32Vec16 gate_vec(gate + n);
94 + vec_op::FP32Vec16 up_vec(up + n);
95 + auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec));
96 + auto silu = gate_vec * sigmoid_vec;
97 + auto gated_output_fp32 = up_vec * silu;
98 + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
99 + gated_output.save(output + n);
100 + }
101 + gate += input_stride;
102 + up += input_stride;
103 + output += output_stride;
104 + }
105 +}
106 +
107 +template <typename scalar_t>
108 +FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
109 + float* __restrict__ input,
110 + scalar_t* __restrict__ output,
111 + const int32_t m, const int32_t n,
112 + const int32_t input_stride,
113 + const int32_t output_stride) {
114 + switch (act) {
115 + case FusedMOEAct::SwigluOAIAndMul:
116 + swigluoai_and_mul(input, output, m, n, input_stride, output_stride);
117 + return;
118 + case FusedMOEAct::SiluAndMul:
119 + silu_and_mul(input, output, m, n, input_stride, output_stride);
120 + return;
121 + default:
122 + TORCH_CHECK(false, "Unsupported act type.");
123 + }
124 +}
125 +
126 +template <typename scalar_t, typename gemm_t>
127 +void prepack_moe_weight_impl(scalar_t* __restrict__ weight_ptr,
128 + scalar_t* __restrict__ packed_weight_ptr,
129 + const int32_t expert_num,
130 + const int32_t output_size,
131 + const int32_t input_size,
132 + const int64_t expert_stride) {
133 +#pragma omp parallel for
134 + for (int32_t e_idx = 0; e_idx < expert_num; ++e_idx) {
135 + gemm_t::pack_weight(weight_ptr + expert_stride * e_idx,
136 + packed_weight_ptr + expert_stride * e_idx, output_size,
137 + input_size);
138 + }
139 +}
140 +
141 +template <typename scalar_t, typename w_t, typename gemm_t>
142 +void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
143 + w_t* __restrict__ w13, w_t* __restrict__ w2,
144 + w_t* __restrict__ w13_bias, w_t* __restrict__ w2_bias,
145 + float* __restrict__ topk_weights,
146 + int32_t* __restrict__ topk_id, FusedMOEAct act_type,
147 + const int32_t token_num, const int32_t expert_num,
148 + const int32_t topk_num, const int32_t input_size_13,
149 + const int32_t output_size_13, const int32_t input_size_2,
150 + const int32_t output_size_2) {
151 + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
152 + constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
153 + constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
154 + constexpr int32_t min_w13_n_tile_size = 2 * gemm_n_tile_size;
155 + static_assert(gemm_n_tile_size % 16 == 0);
156 +
157 + TORCH_CHECK_EQ(output_size_13 % min_w13_n_tile_size, 0);
158 + TORCH_CHECK_EQ(output_size_2 % gemm_n_tile_size, 0);
159 + TORCH_CHECK_EQ(output_size_13 / 2, input_size_2);
160 +
161 + const int32_t thread_num = omp_get_max_threads();
162 +
163 + const int32_t w13_input_buffer_size = cpu_utils::round_up<64>(
164 + gemm_m_tile_size * input_size_13 * sizeof(scalar_t));
165 +
166 + const int32_t w13_n_tile_size = [&]() {
167 + const int64_t cache_size = cpu_utils::get_available_l2_size();
168 + // input buffer + output buffer + weight
169 + const int32_t n_size_cache_limit =
170 + (cache_size - w13_input_buffer_size) /
171 + (gemm_m_tile_size * sizeof(float) + input_size_13 * sizeof(scalar_t));
172 + const int32_t n_size_thread_limit =
173 + output_size_13 / std::max(1, thread_num / topk_num);
174 + const int32_t n_size = cpu_utils::round_down<min_w13_n_tile_size>(
175 + std::min(n_size_cache_limit, n_size_thread_limit));
176 + return std::max(n_size, min_w13_n_tile_size);
177 + }();
178 +
179 + const int32_t w2_input_tile_size = cpu_utils::round_up<64>(
180 + gemm_m_tile_size * input_size_2 * sizeof(scalar_t));
181 +
182 + const int32_t w2_n_tile_size = [&]() {
183 + const int64_t cache_size = cpu_utils::get_available_l2_size();
184 + // input tile + weight
185 + const int32_t n_size_cache_limit =
186 + (cache_size - w2_input_tile_size) / (input_size_2 * sizeof(scalar_t));
187 + const int32_t n_size_thread_limit =
188 + output_size_2 / std::max(1, thread_num / topk_num);
189 + const int32_t n_size = cpu_utils::round_down<gemm_n_tile_size>(
190 + std::min(n_size_cache_limit, n_size_thread_limit));
191 + return std::max(n_size, gemm_n_tile_size);
192 + }();
193 +
194 + // allocate buffers
195 + int32_t common_buffer_offset = 0;
196 + int32_t w13_thread_buffer_offset = 0;
197 + int32_t ws_thread_buffer_offset = 0;
198 +
199 + // common buffers
200 + const int32_t token_num_per_group_buffer_size =
201 + cpu_utils::round_up<64>(expert_num * sizeof(int32_t));
202 + const int32_t token_num_per_group_buffer_offset = common_buffer_offset;
203 + common_buffer_offset += token_num_per_group_buffer_size;
204 +
205 + const int32_t cu_token_num_per_group_buffer_size =
206 + cpu_utils::round_up<64>((expert_num + 1) * sizeof(int32_t));
207 + const int32_t cu_token_num_per_group_buffer_offset = common_buffer_offset;
208 + common_buffer_offset += cu_token_num_per_group_buffer_size;
209 +
210 + const int32_t expand_token_id_buffer_size =
211 + cpu_utils::round_up<64>(token_num * topk_num * sizeof(int32_t));
212 + const int32_t expand_token_id_buffer_offset = common_buffer_offset;
213 + common_buffer_offset += expand_token_id_buffer_size;
214 +
215 + const int32_t expand_token_id_index_buffer_size =
216 + cpu_utils::round_up<64>(token_num * topk_num * sizeof(int32_t));
217 + const int32_t expand_token_id_index_buffer_offset = common_buffer_offset;
218 + common_buffer_offset += expand_token_id_index_buffer_size;
219 +
220 + const int32_t w13_gemm_output_buffer_size = cpu_utils::round_up<64>(
221 + token_num * topk_num * (output_size_13 / 2) * sizeof(scalar_t));
222 + const int32_t w13_gemm_output_buffer_offset = common_buffer_offset;
223 + common_buffer_offset += w13_gemm_output_buffer_size;
224 +
225 + const int32_t w2_gemm_output_buffer_size = cpu_utils::round_up<64>(
226 + token_num * topk_num * output_size_2 * sizeof(float));
227 + const int32_t w2_gemm_output_buffer_offset = common_buffer_offset;
228 + common_buffer_offset += w2_gemm_output_buffer_size;
229 +
230 + // w13 GEMM thread buffers
231 + const int32_t w13_input_buffer_offset = w13_thread_buffer_offset;
232 + w13_thread_buffer_offset += w13_input_buffer_size;
233 +
234 + const int32_t w13_output_buffer_size = cpu_utils::round_up<64>(
235 + gemm_m_tile_size * w13_n_tile_size * sizeof(float));
236 + const int32_t w13_output_buffer_offset = w13_thread_buffer_offset;
237 + w13_thread_buffer_offset += w13_output_buffer_size;
238 +
239 + // Weighted sum thread buffer
240 + const int32_t ws_output_buffer_size =
241 + cpu_utils::round_up<64>(output_size_2 * sizeof(float));
242 + const int32_t ws_output_buffer_offset = ws_thread_buffer_offset;
243 + ws_thread_buffer_offset += ws_output_buffer_size;
244 +
245 + const int32_t buffer_size =
246 + common_buffer_offset +
247 + std::max(w13_thread_buffer_offset, ws_thread_buffer_offset) * thread_num;
248 + cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size);
249 + uint8_t* common_buffer_start =
250 + cpu_utils::ScratchPadManager::get_scratchpad_manager()
251 + ->get_data<uint8_t>();
252 + uint8_t* thread_buffer_start = common_buffer_start + common_buffer_offset;
253 +
254 + int32_t* __restrict__ token_num_per_group_buffer = reinterpret_cast<int32_t*>(
255 + common_buffer_start + token_num_per_group_buffer_offset);
256 + int32_t* __restrict__ cu_token_num_per_group_buffer =
257 + reinterpret_cast<int32_t*>(common_buffer_start +
258 + cu_token_num_per_group_buffer_offset);
259 + int32_t* __restrict__ expand_token_id_buffer = reinterpret_cast<int32_t*>(
260 + common_buffer_start + expand_token_id_buffer_offset);
261 + int32_t* __restrict__ expand_token_id_index_buffer =
262 + reinterpret_cast<int32_t*>(common_buffer_start +
263 + expand_token_id_index_buffer_offset);
264 +
265 + // prepare token-expert mappings
266 + {
267 + std::memset(token_num_per_group_buffer, 0, expert_num * sizeof(int32_t));
268 + for (int32_t i = 0; i < token_num * topk_num; ++i) {
269 + int32_t curr_expert_id = topk_id[i];
270 + ++token_num_per_group_buffer[curr_expert_id];
271 + }
272 +
273 + int32_t token_num_sum = 0;
274 + cu_token_num_per_group_buffer[0] = 0;
275 + int32_t* token_index_buffer = cu_token_num_per_group_buffer + 1;
276 + for (int32_t i = 0; i < expert_num; ++i) {
277 + token_index_buffer[i] = token_num_sum;
278 + token_num_sum += token_num_per_group_buffer[i];
279 + }
280 +
281 + for (int32_t i = 0; i < token_num; ++i) {
282 + int32_t* curr_topk_id = topk_id + i * topk_num;
283 + int32_t* curr_index_buffer = expand_token_id_index_buffer + i * topk_num;
284 + for (int32_t j = 0; j < topk_num; ++j) {
285 + int32_t curr_expert_id = curr_topk_id[j];
286 + int32_t curr_index = token_index_buffer[curr_expert_id];
287 + ++token_index_buffer[curr_expert_id];
288 + expand_token_id_buffer[curr_index] = i;
289 + curr_index_buffer[j] = curr_index;
290 + }
291 + }
292 + }
293 +
294 + // w13 GEMM + act
295 + {
296 + alignas(64) cpu_utils::Counter counter;
297 + cpu_utils::Counter* counter_ptr = &counter;
298 +
299 +#pragma omp parallel for schedule(static, 1)
300 + for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
301 + const int32_t task_num_per_expert =
302 + (output_size_13 + w13_n_tile_size - 1) / w13_n_tile_size;
303 + const int32_t task_num = task_num_per_expert * expert_num;
304 +
305 + uint8_t* __restrict__ thread_buffer =
306 + thread_buffer_start + thread_id * w13_thread_buffer_offset;
307 + scalar_t* __restrict__ w13_input_buffer =
308 + reinterpret_cast<scalar_t*>(thread_buffer + w13_input_buffer_offset);
309 + float* __restrict__ w13_output_buffer =
310 + reinterpret_cast<float*>(thread_buffer + w13_output_buffer_offset);
311 + scalar_t* __restrict__ w13_gemm_output_buffer =
312 + reinterpret_cast<scalar_t*>(common_buffer_start +
313 + w13_gemm_output_buffer_offset);
314 +
315 + gemm_t gemm;
316 +
317 + const int32_t input_size_13_bytes = input_size_13 * sizeof(scalar_t);
318 + const int32_t w13_n_group_stride = 16 * input_size_13;
319 + const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13;
320 +
321 + for (;;) {
322 + int32_t task_id = counter_ptr->acquire_counter();
323 + if (task_id >= task_num) {
324 + break;
325 + }
326 +
327 + const int32_t curr_expert_id = task_id / task_num_per_expert;
328 + const int32_t curr_output_group_id = task_id % task_num_per_expert;
329 + const int32_t curr_token_num =
330 + token_num_per_group_buffer[curr_expert_id];
331 + if (curr_token_num == 0) {
332 + continue;
333 + }
334 +
335 + const int32_t actual_n_tile_size =
336 + std::min(w13_n_tile_size,
337 + output_size_13 - curr_output_group_id * w13_n_tile_size);
338 + const int32_t* __restrict__ curr_expand_token_id_buffer =
339 + expand_token_id_buffer +
340 + cu_token_num_per_group_buffer[curr_expert_id];
341 + scalar_t* __restrict__ curr_w13_gemm_output_buffer =
342 + w13_gemm_output_buffer +
343 + cu_token_num_per_group_buffer[curr_expert_id] *
344 + (output_size_13 / 2) +
345 + curr_output_group_id * w13_n_tile_size / 2;
346 +
347 + w_t* __restrict__ w13_weight_ptr_0 = nullptr;
348 + w_t* __restrict__ w13_weight_ptr_1 = nullptr;
349 + w_t* __restrict__ w13_bias_ptr_0 = nullptr;
350 + w_t* __restrict__ w13_bias_ptr_1 = nullptr;
351 + if (act_type == FusedMOEAct::SwigluOAIAndMul) {
352 + // For SwigluOAIAndMul, up and down weights are interleaved
353 + w13_weight_ptr_0 =
354 + w13 + curr_expert_id * input_size_13 * output_size_13 +
355 + curr_output_group_id * w13_n_tile_size * input_size_13;
356 + w13_weight_ptr_1 =
357 + w13_weight_ptr_0 + actual_n_tile_size / 2 * input_size_13;
358 + if (w13_bias != nullptr) {
359 + w13_bias_ptr_0 = w13_bias + curr_expert_id * output_size_13 +
360 + curr_output_group_id * w13_n_tile_size;
361 + w13_bias_ptr_1 = w13_bias_ptr_0 + actual_n_tile_size / 2;
362 + }
363 + } else {
364 + w13_weight_ptr_0 =
365 + w13 + curr_expert_id * input_size_13 * output_size_13 +
366 + curr_output_group_id * (w13_n_tile_size / 2) * input_size_13;
367 + w13_weight_ptr_1 =
368 + w13_weight_ptr_0 + output_size_13 / 2 * input_size_13;
369 + if (w13_bias != nullptr) {
370 + w13_bias_ptr_0 = w13_bias + curr_expert_id * output_size_13 +
371 + curr_output_group_id * (w13_n_tile_size / 2);
372 + w13_bias_ptr_1 = w13_bias_ptr_0 + output_size_13 / 2;
373 + }
374 + }
375 +
376 + scalar_t* __restrict__ curr_w13_input_buffer = w13_input_buffer;
377 + for (int32_t token_idx = 0; token_idx < curr_token_num;
378 + token_idx += gemm_m_tile_size) {
379 + const int32_t actual_token_num =
380 + std::min(gemm_m_tile_size, curr_token_num - token_idx);
381 + // copy inputs
382 + {
383 + scalar_t* __restrict__ curr_w13_input_buffer_iter =
384 + curr_w13_input_buffer;
385 + for (int32_t i = 0; i < actual_token_num; ++i) {
386 + const int32_t curr_token_id = curr_expand_token_id_buffer[i];
387 + int8_t* __restrict__ curr_input_iter = reinterpret_cast<int8_t*>(
388 + input + curr_token_id * input_size_13);
389 + int8_t* __restrict__ curr_output_iter =
390 + reinterpret_cast<int8_t*>(curr_w13_input_buffer_iter);
391 + int32_t j = 0;
392 + for (; j < input_size_13_bytes - 64; j += 64) {
393 + vec_op::INT8Vec64 vec(curr_input_iter);
394 + vec.save(curr_output_iter);
395 + curr_input_iter += 64;
396 + curr_output_iter += 64;
397 + }
398 + vec_op::INT8Vec64 vec(curr_input_iter);
399 + vec.save(curr_output_iter, input_size_13_bytes - j);
400 +
401 + // update
402 + curr_w13_input_buffer_iter += input_size_13;
403 + }
404 + // update
405 + curr_expand_token_id_buffer += actual_token_num;
406 + }
407 +
408 + // gemm + act
409 + {
410 + scalar_t* __restrict__ w13_weight_ptr_0_iter = w13_weight_ptr_0;
411 + scalar_t* __restrict__ w13_weight_ptr_1_iter = w13_weight_ptr_1;
412 + scalar_t* __restrict__ w13_bias_ptr_0_iter = w13_bias_ptr_0;
413 + scalar_t* __restrict__ w13_bias_ptr_1_iter = w13_bias_ptr_1;
414 + scalar_t* __restrict__ curr_w13_input_buffer_iter =
415 + curr_w13_input_buffer;
416 + float* __restrict__ w13_output_buffer_0_iter = w13_output_buffer;
417 + float* __restrict__ w13_output_buffer_1_iter =
418 + w13_output_buffer + actual_n_tile_size / 2;
419 + for (int32_t i = 0; i < actual_n_tile_size;
420 + i += min_w13_n_tile_size) {
421 + gemm.gemm(curr_w13_input_buffer_iter, w13_weight_ptr_0_iter,
422 + w13_output_buffer_0_iter, actual_token_num,
423 + input_size_13, input_size_13, w13_n_group_stride,
424 + actual_n_tile_size, false);
425 +
426 + if (w13_bias != nullptr) {
427 + cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
428 + w13_output_buffer_0_iter, w13_output_buffer_0_iter,
429 + w13_bias_ptr_0_iter, actual_token_num, actual_n_tile_size,
430 + actual_n_tile_size);
431 + w13_bias_ptr_0_iter += gemm_n_tile_size;
432 + }
433 +
434 + gemm.gemm(curr_w13_input_buffer_iter, w13_weight_ptr_1_iter,
435 + w13_output_buffer_1_iter, actual_token_num,
436 + input_size_13, input_size_13, w13_n_group_stride,
437 + actual_n_tile_size, false);
438 +
439 + if (w13_bias != nullptr) {
440 + cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
441 + w13_output_buffer_1_iter, w13_output_buffer_1_iter,
442 + w13_bias_ptr_1_iter, actual_token_num, actual_n_tile_size,
443 + actual_n_tile_size);
444 + w13_bias_ptr_1_iter += gemm_n_tile_size;
445 + }
446 +
447 + // update
448 + w13_weight_ptr_0_iter += w13_n_tile_stride;
449 + w13_weight_ptr_1_iter += w13_n_tile_stride;
450 + w13_output_buffer_0_iter += gemm_n_tile_size;
451 + w13_output_buffer_1_iter += gemm_n_tile_size;
452 + }
453 +
454 + apply_gated_act(act_type, w13_output_buffer,
455 + curr_w13_gemm_output_buffer, actual_token_num,
456 + actual_n_tile_size, actual_n_tile_size,
457 + output_size_13 / 2);
458 +
459 + // update
460 + curr_w13_gemm_output_buffer +=
461 + gemm_m_tile_size * (output_size_13 / 2);
462 + }
463 + }
464 + }
465 + }
466 + }
467 +
468 + // w2 GEMM
469 + {
470 + alignas(64) cpu_utils::Counter counter;
471 + cpu_utils::Counter* counter_ptr = &counter;
472 +
473 +#pragma omp parallel for schedule(static, 1)
474 + for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
475 + const int32_t task_num_per_expert =
476 + (output_size_2 + w2_n_tile_size - 1) / w2_n_tile_size;
477 + const int32_t task_num = task_num_per_expert * expert_num;
478 + scalar_t* __restrict__ w13_gemm_output_buffer =
479 + reinterpret_cast<scalar_t*>(common_buffer_start +
480 + w13_gemm_output_buffer_offset);
481 + float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>(
482 + common_buffer_start + w2_gemm_output_buffer_offset);
483 +
484 + gemm_t gemm;
485 +
486 + const int32_t w2_n_tile_stride = gemm_n_tile_size * input_size_2;
487 + const int32_t w2_n_group_stride = 16 * input_size_2;
488 +
489 + for (;;) {
490 + int32_t task_id = counter_ptr->acquire_counter();
491 + if (task_id >= task_num) {
492 + break;
493 + }
494 +
495 + const int32_t curr_expert_id = task_id / task_num_per_expert;
496 + const int32_t curr_output_group_id = task_id % task_num_per_expert;
497 + const int32_t curr_token_num =
498 + token_num_per_group_buffer[curr_expert_id];
499 + if (curr_token_num == 0) {
500 + continue;
501 + }
502 +
503 + const int32_t actual_n_tile_size =
504 + std::min(w2_n_tile_size,
505 + output_size_2 - curr_output_group_id * w2_n_tile_size);
506 + scalar_t* __restrict__ curr_w13_gemm_output_buffer =
507 + w13_gemm_output_buffer +
508 + cu_token_num_per_group_buffer[curr_expert_id] * input_size_2;
509 + float* __restrict__ curr_w2_gemm_output_buffer =
510 + w2_gemm_output_buffer +
511 + cu_token_num_per_group_buffer[curr_expert_id] * output_size_2 +
512 + curr_output_group_id * w2_n_tile_size;
513 + scalar_t* __restrict__ w2_weight_ptr =
514 + w2 + curr_expert_id * output_size_2 * input_size_2 +
515 + curr_output_group_id * w2_n_tile_size * input_size_2;
516 + scalar_t* __restrict__ w2_bias_ptr = nullptr;
517 + if (w2_bias != nullptr) {
518 + w2_bias_ptr = w2_bias + curr_expert_id * output_size_2 +
519 + curr_output_group_id * w2_n_tile_size;
520 + }
521 +
522 + for (int32_t token_idx = 0; token_idx < curr_token_num;
523 + token_idx += gemm_m_tile_size) {
524 + const int32_t actual_token_num =
525 + std::min(gemm_m_tile_size, curr_token_num - token_idx);
526 +
527 + scalar_t* __restrict__ w2_weight_ptr_iter = w2_weight_ptr;
528 + scalar_t* __restrict__ w2_bias_ptr_iter = w2_bias_ptr;
529 + float* __restrict__ curr_w2_gemm_output_buffer_iter =
530 + curr_w2_gemm_output_buffer;
531 + for (int32_t i = 0; i < actual_n_tile_size; i += gemm_n_tile_size) {
532 + gemm.gemm(curr_w13_gemm_output_buffer, w2_weight_ptr_iter,
533 + curr_w2_gemm_output_buffer_iter, actual_token_num,
534 + input_size_2, input_size_2, w2_n_group_stride,
535 + output_size_2, false);
536 +
537 + if (w2_bias != nullptr) {
538 + cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
539 + curr_w2_gemm_output_buffer_iter,
540 + curr_w2_gemm_output_buffer_iter, w2_bias_ptr_iter,
541 + actual_token_num, output_size_2, output_size_2);
542 + w2_bias_ptr_iter += gemm_n_tile_size;
543 + }
544 +
545 + w2_weight_ptr_iter += w2_n_tile_stride;
546 + curr_w2_gemm_output_buffer_iter += gemm_n_tile_size;
547 + }
548 +
549 + // update
550 + curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2;
551 + curr_w2_gemm_output_buffer += gemm_m_tile_size * output_size_2;
552 + }
553 + }
554 + }
555 + }
556 +
557 + // weighted sum
558 + {
559 + alignas(64) cpu_utils::Counter counter;
560 + cpu_utils::Counter* counter_ptr = &counter;
561 +
562 +#pragma omp parallel for schedule(static, 1)
563 + for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
564 + const int32_t task_num = token_num;
565 + uint8_t* __restrict__ thread_buffer =
566 + thread_buffer_start + thread_id * ws_thread_buffer_offset;
567 + float* __restrict__ ws_output_buffer =
568 + reinterpret_cast<float*>(thread_buffer + ws_output_buffer_offset);
569 + float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>(
570 + common_buffer_start + w2_gemm_output_buffer_offset);
571 +
572 + for (;;) {
573 + int32_t task_id = counter_ptr->acquire_counter();
574 + if (task_id >= task_num) {
575 + break;
576 + }
577 +
578 + int32_t token_id = task_id;
579 + int32_t* __restrict__ curr_expand_token_id_index_buffer =
580 + expand_token_id_index_buffer + token_id * topk_num;
581 + float* __restrict__ curr_weight = topk_weights + token_id * topk_num;
582 + scalar_t* __restrict__ curr_output_buffer =
583 + output + token_id * output_size_2;
584 +
585 + if (topk_num > 1) {
586 + {
587 + int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
588 + float* __restrict__ w2_output_iter =
589 + w2_gemm_output_buffer + w2_output_idx * output_size_2;
590 + float* __restrict__ ws_output_buffer_iter = ws_output_buffer;
591 + vec_op::FP32Vec16 weight_vec(curr_weight[0]);
592 + for (int32_t i = 0; i < output_size_2; i += 16) {
593 + vec_op::FP32Vec16 vec(w2_output_iter);
594 + vec = vec * weight_vec;
595 + vec.save(ws_output_buffer_iter);
596 +
597 + // update
598 + w2_output_iter += 16;
599 + ws_output_buffer_iter += 16;
600 + }
601 + }
602 +
603 + {
604 + for (int32_t idx = 1; idx < topk_num - 1; ++idx) {
605 + int32_t w2_output_idx = curr_expand_token_id_index_buffer[idx];
606 + float* __restrict__ w2_output_iter =
607 + w2_gemm_output_buffer + w2_output_idx * output_size_2;
608 + float* __restrict__ ws_output_buffer_iter = ws_output_buffer;
609 + vec_op::FP32Vec16 weight_vec(curr_weight[idx]);
610 + for (int32_t i = 0; i < output_size_2; i += 16) {
611 + vec_op::FP32Vec16 vec(w2_output_iter);
612 + vec_op::FP32Vec16 sum(ws_output_buffer_iter);
613 + sum = sum + vec * weight_vec;
614 + sum.save(ws_output_buffer_iter);
615 +
616 + // update
617 + w2_output_iter += 16;
618 + ws_output_buffer_iter += 16;
619 + }
620 + }
621 + }
622 +
623 + {
624 + int32_t idx = topk_num - 1;
625 + int32_t w2_output_idx = curr_expand_token_id_index_buffer[idx];
626 + float* __restrict__ w2_output_iter =
627 + w2_gemm_output_buffer + w2_output_idx * output_size_2;
628 + float* __restrict__ ws_output_buffer_iter = ws_output_buffer;
629 + scalar_t* __restrict__ curr_output_buffer_iter = curr_output_buffer;
630 + vec_op::FP32Vec16 weight_vec(curr_weight[idx]);
631 + for (int32_t i = 0; i < output_size_2; i += 16) {
632 + vec_op::FP32Vec16 vec(w2_output_iter);
633 + vec_op::FP32Vec16 sum(ws_output_buffer_iter);
634 + sum = sum + vec * weight_vec;
635 + scalar_vec_t out_vec(sum);
636 + out_vec.save(curr_output_buffer_iter);
637 +
638 + // update
639 + w2_output_iter += 16;
640 + ws_output_buffer_iter += 16;
641 + curr_output_buffer_iter += 16;
642 + }
643 + }
644 + } else {
645 + int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
646 + float* __restrict__ w2_output_iter =
647 + w2_gemm_output_buffer + w2_output_idx * output_size_2;
648 + scalar_t* __restrict__ curr_output_buffer_iter = curr_output_buffer;
649 + vec_op::FP32Vec16 weight_vec(curr_weight[0]);
650 + for (int32_t i = 0; i < output_size_2; i += 16) {
651 + vec_op::FP32Vec16 vec(w2_output_iter);
652 + vec = vec * weight_vec;
653 + scalar_vec_t out_vec(vec);
654 + out_vec.save(curr_output_buffer_iter);
655 +
656 + // update
657 + w2_output_iter += 16;
658 + curr_output_buffer_iter += 16;
659 + }
660 + }
661 + }
662 + }
663 + }
664 +}
665 +} // namespace
666 +
667 +void prepack_moe_weight(
668 + const torch::Tensor& weight, // [expert_num, output_size, input_size]
669 + torch::Tensor& packed_weight, const std::string& isa) {
670 + TORCH_CHECK(weight.is_contiguous());
671 + const int32_t expert_num = weight.size(0);
672 + const int32_t output_size = weight.size(1);
673 + const int32_t input_size = weight.size(2);
674 + TORCH_CHECK_EQ(output_size % 32, 0);
675 + const int64_t expert_stride = weight.stride(0);
676 + cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
677 +
678 + VLLM_DISPATCH_FLOATING_TYPES(
679 + weight.scalar_type(), "prepack_moe_weight", [&]() {
680 + CPU_ISA_DISPATCH_IMPL(isa_type, [&]() {
681 + scalar_t* weight_ptr = weight.data_ptr<scalar_t>();
682 + scalar_t* packed_weight_ptr = packed_weight.data_ptr<scalar_t>();
683 + prepack_moe_weight_impl<scalar_t, gemm_t>(
684 + weight_ptr, packed_weight_ptr, expert_num, output_size,
685 + input_size, expert_stride);
686 + });
687 + });
688 +}
689 +
690 +void cpu_fused_moe(
691 + torch::Tensor& output, // [token_num, output_size_2]
692 + const torch::Tensor& input, // [token_num, input_size_13]
693 + const torch::Tensor&
694 + w13, // [expert_num, output_size_13, input_size_13], packed
695 + const torch::Tensor&
696 + w2, // [expert_num, output_size_2, input_size_2], packed
697 + const std::optional<torch::Tensor>&
698 + w13_bias, // [expert_num, output_size_13]
699 + const std::optional<torch::Tensor>& w2_bias, // [expert_num, output_size_2]
700 + const torch::Tensor& topk_weights, // [token_num, k], float32
701 + const torch::Tensor& topk_id, // [token_num, k], int32
702 + const std::string& act, const std::string& isa) {
703 + const int32_t token_num = input.size(0);
704 + const int32_t input_size_13 = input.size(1);
705 + const int64_t input_stride = input.stride(0);
706 + TORCH_CHECK_EQ(input_stride, input_size_13);
707 + const int32_t expert_num = w13.size(0);
708 + const int32_t output_size_13 = w13.size(1);
709 + const int32_t input_size_2 = w2.size(2);
710 + const int32_t output_size_2 = w2.size(1);
711 + const int32_t topk_num = topk_id.size(1);
712 + const FusedMOEAct act_type = get_act_type(act);
713 + cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
714 +
715 + VLLM_DISPATCH_FLOATING_TYPES(w13.scalar_type(), "cpu_fused_moe", [&]() {
716 + CPU_ISA_DISPATCH_IMPL(isa_type, [&]() {
717 + fused_moe_impl<scalar_t, scalar_t, gemm_t>(
718 + output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
719 + w13.data_ptr<scalar_t>(), w2.data_ptr<scalar_t>(),
720 + w13_bias.has_value() ? w13_bias->data_ptr<scalar_t>() : nullptr,
721 + w2_bias.has_value() ? w2_bias->data_ptr<scalar_t>() : nullptr,
722 + topk_weights.data_ptr<float>(), topk_id.data_ptr<int32_t>(), act_type,
723 + token_num, expert_num, topk_num, input_size_13, output_size_13,
724 + input_size_2, output_size_2);
725 + });
726 + });
727 +}
csrc/cpu/cpu_types_x86.hpp
+8
@@ -352,6 +352,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
352 explicit FP32Vec16(bool, void* ptr)
353 : reg((__m512)_mm512_stream_load_si512(ptr)) {}
354
355 + // strided load
356 + explicit FP32Vec16(const float* ptr, INT32Vec16 idx)
357 + : reg(_mm512_i32gather_ps(idx.reg, ptr, 4)) {}
358 +
359 explicit FP32Vec16(__m512 data) : reg(data) {}
360
361 // de-pack 4 bit values
@@ -408,6 +412,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
412 return FP32Vec16(_mm512_sub_ps(reg, b.reg));
413 }
414
415 + FP32Vec16 operator-() const {
416 + return FP32Vec16(_mm512_xor_ps(reg, _mm512_set1_ps(-0.0f)));
417 + }
418 +
419 FP32Vec16 operator/(const FP32Vec16& b) const {
420 return FP32Vec16(_mm512_div_ps(reg, b.reg));
421 }
csrc/cpu/cpu_wna16.cpp
+9 -9
@@ -1,6 +1,5 @@
1 -#include "cpu_types.hpp"
2 -#include "scratchpad_manager.h"
3 -#include "utils.hpp"
1 +#include "cpu/cpu_types.hpp"
2 +#include "cpu/utils.hpp"
3
4 #ifdef CPU_CAPABILITY_AMXBF16
5 #include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp"
@@ -158,7 +157,7 @@ void cpu_gemm_wna16_impl(
157 // a simple schedule policy, just to hold more B tiles in L2 and make sure
158 // each thread has tasks
159 const int32_t n_partition_size = [&]() {
161 - const int64_t cache_size = cpu_utils::get_l2_size();
160 + const int64_t cache_size = cpu_utils::get_available_l2_size();
161 int64_t ps_cache_limit = cache_size / (k_size * sizeof(scalar_t));
162 int64_t ps_thread_limit = n_size / thread_num;
163 ps_cache_limit =
@@ -179,8 +178,8 @@ void cpu_gemm_wna16_impl(
178 const int64_t b_buffer_offset = 0;
179 const int64_t c_buffer_offset = b_buffer_size;
180 const int64_t buffer_size = b_buffer_size + c_buffer_size;
182 - DNNLScratchPadManager::get_dnnl_scratchpad_manager()->realloc(buffer_size *
183 - thread_num);
181 + cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size *
182 + thread_num);
183
184 alignas(64) cpu_utils::Counter counter;
185 cpu_utils::Counter* counter_ptr = &counter;
@@ -190,9 +189,10 @@ void cpu_gemm_wna16_impl(
189 scalar_t* __restrict__ b_buffer = nullptr;
190 float* __restrict__ c_buffer = nullptr;
191 {
193 - uint8_t* buffer_ptr = DNNLScratchPadManager::get_dnnl_scratchpad_manager()
194 - ->get_data<uint8_t>() +
195 - thread_id * buffer_size;
192 + uint8_t* buffer_ptr =
193 + cpu_utils::ScratchPadManager::get_scratchpad_manager()
194 + ->get_data<uint8_t>() +
195 + thread_id * buffer_size;
196 b_buffer = reinterpret_cast<scalar_t*>(buffer_ptr + b_buffer_offset);
197 c_buffer = reinterpret_cast<float*>(buffer_ptr + c_buffer_offset);
198 }
csrc/cpu/dnnl_helper.cpp
+6 -6
@@ -4,8 +4,8 @@
4 #include "common/memory_desc.hpp"
5 #include "common/memory.hpp"
6
7 -#include "dnnl_helper.h"
8 -#include "scratchpad_manager.h"
7 +#include "cpu/utils.hpp"
8 +#include "cpu/dnnl_helper.h"
9
10 static dnnl::engine& default_engine() {
11 static dnnl::engine engine(dnnl::engine::kind::cpu, 0);
@@ -274,7 +274,7 @@ void W8A8MatMulPrimitiveHandler::execute(ExecArgs& args) {
274
275 auto&& [scratchpad_storage, scratchpad_mem_desc] = get_runtime_memory_ptr(5);
276 scratchpad_storage->set_data_handle(
277 - DNNLScratchPadManager::get_dnnl_scratchpad_manager()->get_data<void>());
277 + cpu_utils::ScratchPadManager::get_scratchpad_manager()->get_data<void>());
278
279 matmul.execute(default_stream(), memory_cache_);
280 default_stream().wait();
@@ -294,7 +294,7 @@ dnnl::matmul W8A8MatMulPrimitiveHandler::get_matmul_cache(
294
295 return m_size_cache_->get_or_create(key, [&]() {
296 dnnl::matmul::primitive_desc desc = this->create_primitive_desc(key, false);
297 - auto manager = DNNLScratchPadManager::get_dnnl_scratchpad_manager();
297 + auto manager = cpu_utils::ScratchPadManager::get_scratchpad_manager();
298 manager->realloc(desc.scratchpad_desc().get_size());
299 return dnnl::matmul(desc);
300 });
@@ -470,7 +470,7 @@ void MatMulPrimitiveHandler::execute(ExecArgs& args) {
470
471 auto&& [scratchpad_storage, scratchpad_mem_desc] = get_runtime_memory_ptr(3);
472 scratchpad_storage->set_data_handle(
473 - DNNLScratchPadManager::get_dnnl_scratchpad_manager()->get_data<void>());
473 + cpu_utils::ScratchPadManager::get_scratchpad_manager()->get_data<void>());
474
475 matmul.execute(default_stream(), memory_cache_);
476 default_stream().wait();
@@ -486,7 +486,7 @@ dnnl::matmul MatMulPrimitiveHandler::get_matmul_cache(
486 }
487 return m_size_cache_->get_or_create(key, [&]() {
488 dnnl::matmul::primitive_desc desc = this->create_primitive_desc(key, false);
489 - auto manager = DNNLScratchPadManager::get_dnnl_scratchpad_manager();
489 + auto manager = cpu_utils::ScratchPadManager::get_scratchpad_manager();
490 manager->realloc(desc.scratchpad_desc().get_size());
491 return dnnl::matmul(desc);
492 });
csrc/cpu/micro_gemm/cpu_micro_gemm_amx.hpp
+33
@@ -235,6 +235,39 @@ class MicroGemm<cpu_utils::ISA::AMX, scalar_t> {
235 }
236 }
237
238 + static void pack_weight(const scalar_t* __restrict__ weight,
239 + scalar_t* __restrict__ packed_weight,
240 + const int32_t output_size, const int32_t input_size) {
241 + constexpr int32_t elem_num_per_group = 4 / sizeof(scalar_t);
242 + TORCH_CHECK_EQ(output_size % 16, 0);
243 + TORCH_CHECK_EQ(input_size % (16 * elem_num_per_group), 0);
244 +
245 + const int32_t output_group_num = output_size / 16;
246 + const int32_t input_32b_num = input_size / elem_num_per_group;
247 + for (int32_t output_group_idx = 0; output_group_idx < output_group_num;
248 + ++output_group_idx) {
249 + const int32_t* __restrict__ weight_32b =
250 + reinterpret_cast<const int32_t*>(weight);
251 + int32_t* __restrict__ packed_weight_32b =
252 + reinterpret_cast<int32_t*>(packed_weight);
253 + for (int32_t output_idx = 0; output_idx < 16; ++output_idx) {
254 + for (int32_t weight_offset = 0, packed_offset = 0;
255 + weight_offset < input_32b_num;
256 + ++weight_offset, packed_offset += 16) {
257 + packed_weight_32b[packed_offset] = weight_32b[weight_offset];
258 + }
259 +
260 + // update
261 + weight_32b += input_32b_num;
262 + packed_weight_32b += 1;
263 + }
264 +
265 + // update
266 + weight += 16 * input_size;
267 + packed_weight += 16 * input_size;
268 + }
269 + }
270 +
271 private:
272 alignas(64) __tilecfg amx_tile_config_;
273 int32_t curr_m_;
csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp
+38
@@ -13,6 +13,9 @@ namespace cpu_micro_gemm {
13 #define CPU_MICRO_GEMM_PARAMS \
14 a_ptr, b_ptr, c_ptr, m, k, lda, b_n_group_stride, ldc, accum_c
15
16 +// Note: weights for MicroGemm should be packed as (output_size / 16) contiguous
17 +// blocks, means the logical shape of blocks is [16, input_size]. And the actual
18 +// layout of blocks can be ISA-specific.
19 template <cpu_utils::ISA isa, typename scalar_t>
20 class MicroGemm {
21 public:
@@ -86,6 +89,41 @@ FORCE_INLINE void bias_epilogue(float* __restrict__ c_ptr,
89 curr_d += ldd;
90 }
91 }
92 +
93 +template <int32_t n_size, typename scalar_t>
94 +FORCE_INLINE void add_bias_epilogue(float* c_ptr, float* d_ptr,
95 + scalar_t* __restrict__ bias_ptr,
96 + const int32_t m, const int64_t ldc,
97 + const int64_t ldd) {
98 + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
99 + static_assert(n_size % 16 == 0);
100 + constexpr int32_t n_group_num = n_size / 16;
101 + static_assert(n_group_num <= 16);
102 +
103 + vec_op::FP32Vec16 bias_vecs[n_group_num];
104 + scalar_t* __restrict__ curr_bias = bias_ptr;
105 + vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t i) {
106 + scalar_vec_t vec(curr_bias);
107 + bias_vecs[i] = vec_op::FP32Vec16(vec);
108 + curr_bias += 16;
109 + });
110 +
111 + float* curr_c = c_ptr;
112 + float* curr_d = d_ptr;
113 + for (int32_t i = 0; i < m; ++i) {
114 + float* curr_c_iter = curr_c;
115 + float* curr_d_iter = curr_d;
116 + vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t n_g_idx) {
117 + vec_op::FP32Vec16 c_vec_fp32(curr_c_iter);
118 + c_vec_fp32 = c_vec_fp32 + bias_vecs[n_g_idx];
119 + c_vec_fp32.save(curr_d_iter);
120 + curr_c_iter += 16;
121 + curr_d_iter += 16;
122 + });
123 + curr_c += ldc;
124 + curr_d += ldd;
125 + }
126 +}
127 } // namespace cpu_micro_gemm
128
129 #endif
csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp
+19
@@ -109,6 +109,25 @@ class MicroGemm<cpu_utils::ISA::VEC, scalar_t> {
109 void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
110 TileGemm82<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
111 }
112 +
113 + // Note: pack contiguous weight [output_size, input_size] as contiguous
114 + // packed weight [output_size / 16, input_size, 16]
115 + static void pack_weight(const scalar_t* __restrict__ weight,
116 + scalar_t* __restrict__ packed_weight,
117 + const int32_t output_size, const int32_t input_size) {
118 + TORCH_CHECK_EQ(output_size % 16, 0);
119 + for (int32_t o_idx = 0; o_idx < output_size; ++o_idx) {
120 + const scalar_t* __restrict__ curr_weight = weight + o_idx * input_size;
121 + scalar_t* __restrict__ curr_packed_weight =
122 + packed_weight + (o_idx / 16) * (16 * input_size) + o_idx % 16;
123 + for (int32_t i_idx = 0; i_idx < input_size; ++i_idx) {
124 + *curr_packed_weight = *curr_weight;
125 +
126 + curr_packed_weight += 16;
127 + ++curr_weight;
128 + }
129 + }
130 + }
131 };
132 } // namespace cpu_micro_gemm
133
csrc/cpu/scratchpad_manager.cpp deleted
-23
@@ -1,23 +0,0 @@
1 -#include <cstdlib>
2 -
3 -#include "scratchpad_manager.h"
4 -
5 -DNNLScratchPadManager::DNNLScratchPadManager() : size_(0), ptr_(nullptr) {
6 - this->realloc(allocation_unit * 128);
7 -}
8 -
9 -void DNNLScratchPadManager::realloc(size_t new_size) {
10 - new_size = round(new_size);
11 - if (new_size > size_) {
12 - if (ptr_ != nullptr) {
13 - std::free(ptr_);
14 - }
15 - ptr_ = std::aligned_alloc(64, new_size);
16 - size_ = new_size;
17 - }
18 -}
19 -
20 -DNNLScratchPadManager* DNNLScratchPadManager::get_dnnl_scratchpad_manager() {
21 - static DNNLScratchPadManager manager;
22 - return &manager;
23 -}
csrc/cpu/scratchpad_manager.h deleted
-31
@@ -1,31 +0,0 @@
1 -#ifndef SCRATCHPAD_MANAGER_H
2 -#define SCRATCHPAD_MANAGER_H
3 -
4 -#include <cstddef>
5 -#include <cstdio>
6 -
7 -class DNNLScratchPadManager {
8 - public:
9 - static constexpr size_t allocation_unit = 4 * 1024; // 4KB
10 -
11 - static DNNLScratchPadManager* get_dnnl_scratchpad_manager();
12 -
13 - DNNLScratchPadManager();
14 -
15 - template <typename T>
16 - T* get_data() {
17 - return reinterpret_cast<T*>(ptr_);
18 - }
19 -
20 - static size_t round(size_t size) {
21 - return ((size + allocation_unit - 1) / allocation_unit) * allocation_unit;
22 - }
23 -
24 - void realloc(size_t new_size);
25 -
26 - private:
27 - size_t size_;
28 - void* ptr_;
29 -};
30 -
31 -#endif
csrc/cpu/torch_bindings.cpp
+24
@@ -110,6 +110,17 @@ void cpu_gemm_wna16(const torch::Tensor& input, const torch::Tensor& q_weight,
110 const std::optional<torch::Tensor>& bias,
111 const int64_t pack_factor, const std::string& isa_hint);
112
113 +void prepack_moe_weight(const torch::Tensor& weight,
114 + torch::Tensor& packed_weight, const std::string& isa);
115 +
116 +void cpu_fused_moe(torch::Tensor& output, const torch::Tensor& input,
117 + const torch::Tensor& w13, const torch::Tensor& w2,
118 + const std::optional<torch::Tensor>& w13_bias,
119 + const std::optional<torch::Tensor>& w2_bias,
120 + const torch::Tensor& topk_weights,
121 + const torch::Tensor& topk_id, const std::string& act,
122 + const std::string& isa);
123 +
124 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
125 // vLLM custom ops
126
@@ -296,6 +307,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
307 "pack_factor, str isa_hint) -> ()");
308 ops.impl("cpu_gemm_wna16", torch::kCPU, &cpu_gemm_wna16);
309 #endif
310 +
311 + // fused moe
312 +#if defined(__AVX512F__)
313 + ops.def(
314 + "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
315 + "-> ()");
316 + ops.impl("prepack_moe_weight", torch::kCPU, &prepack_moe_weight);
317 + ops.def(
318 + "cpu_fused_moe(Tensor(a0!) output, Tensor input, Tensor w13, Tensor w2, "
319 + "Tensor? w13_bias, Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, "
320 + "str act, str isa) -> ()");
321 + ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
322 +#endif
323 }
324
325 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _utils), utils) {
csrc/cpu/utils.cpp
+23 -1
@@ -10,7 +10,7 @@
10 #define gettid() syscall(SYS_gettid)
11 #endif
12
13 -#include "cpu_types.hpp"
13 +#include "cpu/utils.hpp"
14
15 #ifdef VLLM_NUMA_DISABLED
16 std::string init_cpu_threads_env(const std::string& cpu_ids) {
@@ -138,4 +138,26 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {
138
139 return ss.str();
140 }
141 +
142 +namespace cpu_utils {
143 +ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {
144 + this->realloc(allocation_unit * 128);
145 +}
146 +
147 +void ScratchPadManager::realloc(size_t new_size) {
148 + new_size = round(new_size);
149 + if (new_size > size_) {
150 + if (ptr_ != nullptr) {
151 + std::free(ptr_);
152 + }
153 + ptr_ = std::aligned_alloc(64, new_size);
154 + size_ = new_size;
155 + }
156 +}
157 +
158 +ScratchPadManager* ScratchPadManager::get_scratchpad_manager() {
159 + static ScratchPadManager manager;
160 + return &manager;
161 +}
162 +} // namespace cpu_utils
163 #endif
csrc/cpu/utils.hpp
+67 -22
@@ -2,19 +2,24 @@
2 #define UTILS_HPP
3
4 #include <atomic>
5 -#include <cassert>
6 -#include <cstdint>
5 #include <unistd.h>
6 +#include <ATen/cpu/Utils.h>
7
9 -#if defined(__APPLE__)
10 - #include <sys/sysctl.h>
11 -#endif
12 -
13 -#include "cpu_types.hpp"
8 +#include "cpu/cpu_types.hpp"
9
10 namespace cpu_utils {
11 enum class ISA { AMX, VEC };
12
13 +inline ISA get_isa(const std::string& isa) {
14 + if (isa == "amx") {
15 + return ISA::AMX;
16 + } else if (isa == "vec") {
17 + return ISA::VEC;
18 + } else {
19 + TORCH_CHECK(false, "Invalid isa type: " + isa);
20 + }
21 +}
22 +
23 template <typename T>
24 struct VecTypeTrait {
25 using vec_t = void;
@@ -48,26 +53,66 @@ struct Counter {
53 int64_t acquire_counter() { return counter++; }
54 };
55
51 -inline int64_t get_l2_size() {
56 +inline int64_t get_available_l2_size() {
57 static int64_t size = []() {
53 -#if defined(__APPLE__)
54 - // macOS doesn't have _SC_LEVEL2_CACHE_SIZE. Use sysctlbyname.
55 - int64_t l2_cache_size = 0;
56 - size_t len = sizeof(l2_cache_size);
57 - if (sysctlbyname("hw.l2cachesize", &l2_cache_size, &len, NULL, 0) == 0 &&
58 - l2_cache_size > 0) {
59 - return l2_cache_size >> 1; // use 50% of L2 cache
60 - }
61 - // Fallback if sysctlbyname fails
62 - return 128LL * 1024 >> 1; // use 50% of 128KB
63 -#else
64 - long l2_cache_size = sysconf(_SC_LEVEL2_CACHE_SIZE);
65 - assert(l2_cache_size != -1);
58 + const uint32_t l2_cache_size = at::cpu::L2_cache_size();
59 return l2_cache_size >> 1; // use 50% of L2 cache
67 -#endif
60 }();
61 return size;
62 }
63 +
64 +template <int32_t alignment_v, typename T>
65 +inline T round_up(T size) {
66 + T alignment = alignment_v;
67 + return (((size + alignment - 1) / alignment) * alignment);
68 +}
69 +
70 +template <int32_t alignment_v, typename T>
71 +inline T round_down(T size) {
72 + T alignment = alignment_v;
73 + return (size / alignment) * alignment;
74 +}
75 +
76 +template <typename T>
77 +inline void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
78 + int32_t stride) {
79 + std::stringstream ss;
80 + ss << std::fixed << std::setprecision(5) << name << ": [\n";
81 + auto* curr_logits_buffer = ptr;
82 + for (int32_t m = 0; m < row; ++m) {
83 + for (int32_t n = 0; n < col; ++n) {
84 + ss << curr_logits_buffer[n] << ", ";
85 + }
86 + ss << "\n";
87 + curr_logits_buffer += stride;
88 + }
89 + ss << "]\n";
90 + std::printf("%s", ss.str().c_str());
91 +}
92 +
93 +class ScratchPadManager {
94 + public:
95 + static constexpr size_t allocation_unit = 4 * 1024; // 4KB
96 +
97 + static ScratchPadManager* get_scratchpad_manager();
98 +
99 + ScratchPadManager();
100 +
101 + template <typename T>
102 + T* get_data() {
103 + return reinterpret_cast<T*>(ptr_);
104 + }
105 +
106 + static size_t round(size_t size) {
107 + return ((size + allocation_unit - 1) / allocation_unit) * allocation_unit;
108 + }
109 +
110 + void realloc(size_t new_size);
111 +
112 + private:
113 + size_t size_;
114 + void* ptr_;
115 +};
116 } // namespace cpu_utils
117
118 #endif
docker/Dockerfile.cpu
+3 -1
@@ -147,7 +147,9 @@ WORKDIR /workspace/vllm
147
148 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
149 --mount=type=cache,target=/var/lib/apt,sharing=locked \
150 - apt-get install -y --no-install-recommends vim numactl xz-utils
150 + apt-get install -y --no-install-recommends vim numactl xz-utils make clangd-14
151 +
152 +RUN ln -s /usr/bin/clangd-14 /usr/bin/clangd
153
154 # install development dependencies (for testing)
155 RUN --mount=type=cache,target=/root/.cache/uv \
requirements/cpu-build.txt
+1 -1
@@ -1,7 +1,7 @@
1 cmake>=3.26.1
2 ninja
3 packaging>=24.2
4 -setuptools>=77.0.3,<81.0.0
4 +setuptools==77.0.3 # this version can reuse CMake build dir
5 setuptools-scm>=8
6 torch==2.9.1+cpu; platform_machine == "x86_64" or platform_machine == "s390x"
7 torch==2.9.1; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "aarch64"
requirements/cpu.txt
+2
@@ -1,6 +1,8 @@
1 # Common dependencies
2 -r common.txt
3
4 +setuptools==77.0.3 # this version can reuse CMake build dir
5 +
6 numba == 0.61.2; platform_machine != "s390x" # Required for N-gram speculative decoding
7
8 # Dependencies for CPUs
tests/kernels/moe/test_cpu_fused_moe.py new
+172
@@ -0,0 +1,172 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from tests.kernels.allclose_default import get_default_atol, get_default_rtol
8 +from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight
9 +from vllm.model_executor.layers.activation import SiluAndMul, SwigluOAIAndMul
10 +from vllm.platforms import current_platform
11 +
12 +if not current_platform.is_cpu():
13 + pytest.skip("skipping CPU-only tests", allow_module_level=True)
14 +
15 +EXPERT_NUM = [
16 + 8,
17 +]
18 +HIDDEN_DIM = [128, 2880]
19 +INTERMEDIATE_DIM = [128, 2880]
20 +BATCH_SIZE = [1, 64, 256]
21 +ACT = ["silu", "swigluoai"]
22 +USE_BIAS = [True, False]
23 +ISA = ["amx", "vec"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]
24 +DTYPE = [torch.bfloat16]
25 +
26 +_CPU_MOE_ACT = {
27 + "silu": SiluAndMul(),
28 + "swigluoai": SwigluOAIAndMul(),
29 +}
30 +
31 +
32 +def ref_fused_moe(
33 + input: torch.Tensor,
34 + w13: torch.Tensor,
35 + w2: torch.Tensor,
36 + w13_bias: torch.Tensor | None,
37 + w2_bias: torch.Tensor | None,
38 + topk_weights: torch.Tensor,
39 + topk_ids: torch.Tensor,
40 + activation: str,
41 +) -> torch.Tensor:
42 + len_experts = w13.size(0)
43 +
44 + cnts = topk_ids.new_zeros((topk_ids.shape[0], len_experts))
45 + cnts.scatter_(1, topk_ids.to(torch.int64), 1)
46 + tokens_per_expert = cnts.sum(dim=0)
47 + idxs = topk_ids.view(-1).argsort()
48 +
49 + sorted_tokens = input[idxs // topk_ids.shape[1]]
50 + tokens_per_expert = tokens_per_expert.cpu().numpy()
51 +
52 + outputs = []
53 + start_idx = 0
54 +
55 + for i, num_tokens in enumerate(tokens_per_expert):
56 + end_idx = start_idx + num_tokens
57 + if num_tokens == 0:
58 + continue
59 + tokens_for_this_expert = sorted_tokens[start_idx:end_idx].float()
60 + curr_w13 = w13[i].float()
61 + curr_w2 = w2[i].float()
62 +
63 + curr_w13_bias = None
64 + if w13_bias is not None:
65 + curr_w13_bias = w13_bias[i].float()
66 +
67 + curr_w2_bias = None
68 + if w2_bias is not None:
69 + curr_w2_bias = w2_bias[i].float()
70 +
71 + gate_up = torch.nn.functional.linear(
72 + tokens_for_this_expert, curr_w13, curr_w13_bias
73 + )
74 + # Note: to simulate the kernel implementation
75 + gate_up = (
76 + _CPU_MOE_ACT[activation]
77 + .forward_native(gate_up)
78 + .to(dtype=input.dtype)
79 + .float()
80 + )
81 + expert_out = torch.nn.functional.linear(gate_up, curr_w2, curr_w2_bias)
82 +
83 + outputs.append(expert_out)
84 + start_idx = end_idx
85 +
86 + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
87 + new_x = torch.empty_like(outs)
88 +
89 + new_x[idxs] = outs
90 + final_out = (
91 + new_x.view(*topk_ids.shape, -1)
92 + .mul_(topk_weights.unsqueeze(dim=-1))
93 + .sum(dim=1)
94 + .type(input.dtype)
95 + )
96 + return final_out
97 +
98 +
99 +@pytest.mark.parametrize("batch_size", BATCH_SIZE)
100 +@pytest.mark.parametrize("expert_num", EXPERT_NUM)
101 +@pytest.mark.parametrize("hidden_size", HIDDEN_DIM)
102 +@pytest.mark.parametrize("intermediate_size", INTERMEDIATE_DIM)
103 +@pytest.mark.parametrize("use_bias", USE_BIAS)
104 +@pytest.mark.parametrize("dtype", DTYPE)
105 +@pytest.mark.parametrize("act", ACT)
106 +@pytest.mark.parametrize("isa", ISA)
107 +def test_cpu_fused_moe(
108 + batch_size: int,
109 + expert_num: int,
110 + hidden_size: int,
111 + intermediate_size: int,
112 + use_bias: bool,
113 + dtype: torch.dtype,
114 + act: str,
115 + isa: str,
116 +):
117 + current_platform.seed_everything(0)
118 +
119 + topk_num = max(expert_num // 2, 1)
120 + up_dim = 2 * intermediate_size
121 +
122 + input = torch.randn((batch_size, hidden_size), dtype=dtype) / (
123 + 0.5 * hidden_size**0.5
124 + )
125 + w13 = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / (
126 + 0.5 * hidden_size**0.5
127 + )
128 + w2 = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / (
129 + 0.5 * intermediate_size**0.5
130 + )
131 + router_logits = torch.randn((batch_size, expert_num), dtype=dtype)
132 + w13_bias = None
133 + w2_bias = None
134 + if use_bias:
135 + w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5)
136 + w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / (
137 + 0.5 * hidden_size**0.5
138 + )
139 + score = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
140 + topk_weight, topk_ids = torch.topk(score, topk_num)
141 + topk_ids = topk_ids.to(torch.int32)
142 +
143 + ref_output = ref_fused_moe(
144 + input,
145 + w13,
146 + w2,
147 + w13_bias,
148 + w2_bias,
149 + topk_weight,
150 + topk_ids,
151 + act,
152 + )
153 +
154 + packed_w13 = cpu_prepack_moe_weight(w13, isa)
155 + packed_w2 = cpu_prepack_moe_weight(w2, isa)
156 + output = cpu_fused_moe(
157 + input,
158 + packed_w13,
159 + packed_w2,
160 + w13_bias,
161 + w2_bias,
162 + topk_weight,
163 + topk_ids,
164 + act,
165 + isa,
166 + )
167 +
168 + atol, rtol = get_default_atol(output), get_default_rtol(output)
169 + (
170 + torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
171 + f"{torch.max(torch.abs(output - ref_output))}",
172 + )
vllm/_custom_ops.py
+36
@@ -2919,6 +2919,42 @@ def cpu_gemm_wna16(
2919 return output
2920
2921
2922 +def cpu_prepack_moe_weight(
2923 + weight: torch.Tensor,
2924 + isa: str,
2925 +) -> torch.Tensor:
2926 + output = torch.empty_like(weight)
2927 + torch.ops._C.prepack_moe_weight(weight, output, isa)
2928 + return output
2929 +
2930 +
2931 +def cpu_fused_moe(
2932 + input: torch.Tensor,
2933 + w13: torch.Tensor,
2934 + w2: torch.Tensor,
2935 + w13_bias: torch.Tensor | None,
2936 + w2_bias: torch.Tensor | None,
2937 + topk_weights: torch.Tensor,
2938 + topk_ids: torch.Tensor,
2939 + act: str,
2940 + isa: str,
2941 +) -> torch.Tensor:
2942 + output = torch.empty_like(input)
2943 + torch.ops._C.cpu_fused_moe(
2944 + output,
2945 + input,
2946 + w13,
2947 + w2,
2948 + w13_bias,
2949 + w2_bias,
2950 + topk_weights,
2951 + topk_ids,
2952 + act,
2953 + isa,
2954 + )
2955 + return output
2956 +
2957 +
2958 if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"):
2959
2960 @register_fake("_qutlass_C::matmul_mxf4_bf16_tn")
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+205 -71
@@ -1,12 +1,22 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import weakref
4 from collections.abc import Callable
5
6 import torch
7 from torch.nn import functional as F
8
9 from vllm import _custom_ops as ops
10 +from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight
11 from vllm.model_executor.layers.activation import SiluAndMul, SwigluOAIAndMul
12 +from vllm.model_executor.layers.quantization.utils.layer_utils import replace_parameter
13 +from vllm.utils.torch_utils import direct_register_custom_op
14 +
15 +_CPU_MOE_LAYER_CACHE = {}
16 +_CPU_MOE_ACT = {
17 + "silu": SiluAndMul(),
18 + "swigluoai": SwigluOAIAndMul(),
19 +}
20
21
22 def grouped_topk(
@@ -174,8 +184,105 @@ class SGLFusedMOE:
184
185 class CPUFusedMOE:
186 def __init__(self, layer: torch.nn.Module) -> None:
177 - use_onednn_mm = ops._supports_onednn and ops.is_onednn_acl_supported()
187 + use_grouped_gemm, isa = self.check_grouped_gemm(layer)
188 + self.isa = isa
189 + if use_grouped_gemm:
190 + self.forward_method = self.forward_grouped_gemm
191 + self.init_moe_grouped_gemm(layer=layer)
192 + else:
193 + self.forward_method = self.forward_torch
194 + self.init_moe_torch(layer=layer)
195 +
196 + def __call__(
197 + self,
198 + layer: torch.nn.Module,
199 + x: torch.Tensor,
200 + use_grouped_topk: bool,
201 + top_k: int,
202 + router_logits: torch.Tensor,
203 + renormalize: bool,
204 + topk_group: int | None = None,
205 + num_expert_group: int | None = None,
206 + global_num_experts: int = -1,
207 + expert_map: torch.Tensor | None = None,
208 + custom_routing_function: Callable | None = None,
209 + scoring_func: str = "softmax",
210 + routed_scaling_factor: float = 1.0,
211 + e_score_correction_bias: torch.Tensor | None = None,
212 + apply_router_weight_on_input: bool = False,
213 + activation: str = "silu",
214 + ) -> torch.Tensor:
215 + assert activation in _CPU_MOE_ACT, f"{activation} is not supported."
216 + assert not apply_router_weight_on_input
217 +
218 + topk_weights, topk_ids = select_experts(
219 + hidden_states=x,
220 + router_logits=router_logits,
221 + use_grouped_topk=use_grouped_topk,
222 + top_k=top_k,
223 + renormalize=renormalize,
224 + topk_group=topk_group,
225 + num_expert_group=num_expert_group,
226 + custom_routing_function=custom_routing_function,
227 + scoring_func=scoring_func,
228 + routed_scaling_factor=routed_scaling_factor,
229 + e_score_correction_bias=e_score_correction_bias,
230 + )
231 +
232 + return self.forward_method(
233 + layer,
234 + x,
235 + topk_weights,
236 + topk_ids,
237 + activation,
238 + global_num_experts,
239 + )
240 +
241 + def check_grouped_gemm(
242 + self,
243 + layer: torch.nn.Module,
244 + ) -> tuple[bool, str]:
245 + if not hasattr(torch.ops._C, "prepack_moe_weight"):
246 + return False, "none"
247 +
248 + dtype = layer.w13_weight.dtype
249 + w13_input_size = layer.w13_weight.size(2)
250 + w13_output_size = layer.w13_weight.size(1)
251 + w2_input_size = layer.w2_weight.size(2)
252 + w2_output_size = layer.w2_weight.size(1)
253 +
254 + if not (w13_output_size % 32 == 0 and w2_output_size % 32 == 0):
255 + return False, "none"
256 +
257 + supports_amx = torch._C._cpu._is_amx_tile_supported()
258 +
259 + if (
260 + supports_amx
261 + and dtype == torch.bfloat16
262 + and w13_input_size % 32 == 0
263 + and w2_input_size % 32 == 0
264 + ):
265 + return True, "amx"
266 +
267 + if supports_amx:
268 + return False, "none"
269 +
270 + return True, "vec"
271 +
272 + def init_moe_grouped_gemm(
273 + self,
274 + layer: torch.nn.Module,
275 + ) -> None:
276 + new_w13 = cpu_prepack_moe_weight(layer.w13_weight, self.isa)
277 + replace_parameter(layer, "w13_weight", new_w13)
278 + new_w2 = cpu_prepack_moe_weight(layer.w2_weight, self.isa)
279 + replace_parameter(layer, "w2_weight", new_w2)
280
281 + def init_moe_torch(
282 + self,
283 + layer: torch.nn.Module,
284 + ) -> None:
285 + use_onednn_mm = ops._supports_onednn and ops.is_onednn_acl_supported()
286 num_experts = layer.w13_weight.size(0)
287 has_w13_bias = hasattr(layer, "w13_bias")
288 has_w2_bias = hasattr(layer, "w2_bias")
@@ -208,85 +315,112 @@ class CPUFusedMOE:
315 layer.down_linear.append(
316 lambda x, w=layer_w2_weight, b=layer_w2_bias: F.linear(x, w, b)
317 )
318 +
319 if use_onednn_mm: # remove weight
320 layer.w13_weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
321 layer.w2_weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
322
215 - self.act_to_impl = {
216 - "silu": SiluAndMul(),
217 - "swigluoai": SwigluOAIAndMul(),
218 - }
323 + _CPU_MOE_LAYER_CACHE[id(layer)] = weakref.ref(layer)
324
220 - def __call__(
325 + def forward_grouped_gemm(
326 self,
327 layer: torch.nn.Module,
223 - x: torch.Tensor,
224 - use_grouped_topk: bool,
225 - top_k: int,
226 - router_logits: torch.Tensor,
227 - renormalize: bool,
228 - topk_group: int | None = None,
229 - num_expert_group: int | None = None,
328 + input: torch.Tensor,
329 + topk_weights: torch.Tensor,
330 + topk_ids: torch.Tensor,
331 + activation: str,
332 global_num_experts: int = -1,
231 - expert_map: torch.Tensor | None = None,
232 - custom_routing_function: Callable | None = None,
233 - scoring_func: str = "softmax",
234 - routed_scaling_factor: float = 1.0,
235 - e_score_correction_bias: torch.Tensor | None = None,
236 - apply_router_weight_on_input: bool = False,
237 - activation: str = "silu",
333 ) -> torch.Tensor:
239 - assert activation in self.act_to_impl, f"{activation} is not supported."
240 - assert not apply_router_weight_on_input
241 - topk_weights, topk_ids = select_experts(
242 - hidden_states=x,
243 - router_logits=router_logits,
244 - use_grouped_topk=use_grouped_topk,
245 - top_k=top_k,
246 - renormalize=renormalize,
247 - topk_group=topk_group,
248 - num_expert_group=num_expert_group,
249 - custom_routing_function=custom_routing_function,
250 - scoring_func=scoring_func,
251 - routed_scaling_factor=routed_scaling_factor,
252 - e_score_correction_bias=e_score_correction_bias,
334 + output = cpu_fused_moe(
335 + input,
336 + layer.w13_weight,
337 + layer.w2_weight,
338 + getattr(layer, "w13_bias", None),
339 + getattr(layer, "w2_bias", None),
340 + topk_weights,
341 + topk_ids,
342 + activation,
343 + self.isa,
344 )
345 + return output
346
255 - # Ref code from https://github.com/sgl-project/sglang/blob/716e682721397df103f347d22da8bd46c6016dab/python/sglang/srt/layers/moe/fused_moe_native.py#L53
256 - len_experts = global_num_experts
257 -
258 - cnts = topk_ids.new_zeros((topk_ids.shape[0], len_experts))
259 - cnts.scatter_(1, topk_ids.to(torch.int64), 1)
260 - tokens_per_expert = cnts.sum(dim=0)
261 - idxs = topk_ids.view(-1).argsort()
262 -
263 - sorted_tokens = x[idxs // topk_ids.shape[1]]
264 - tokens_per_expert = tokens_per_expert.cpu().numpy()
265 -
266 - outputs = []
267 - start_idx = 0
268 -
269 - for i, num_tokens in enumerate(tokens_per_expert):
270 - end_idx = start_idx + num_tokens
271 - if num_tokens == 0:
272 - continue
273 - tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
274 -
275 - gate_up = layer.gate_up_linear[i](tokens_for_this_expert)
276 - gate_up = self.act_to_impl[activation].forward_native(gate_up)
277 - expert_out = layer.down_linear[i](gate_up)
278 - outputs.append(expert_out)
279 - start_idx = end_idx
280 -
281 - outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
282 - new_x = torch.empty_like(outs)
283 -
284 - new_x[idxs] = outs
285 - final_out = (
286 - new_x.view(*topk_ids.shape, -1)
287 - .type(topk_weights.dtype)
288 - .mul_(topk_weights.unsqueeze(dim=-1))
289 - .sum(dim=1)
290 - .type(new_x.dtype)
347 + def forward_torch(
348 + self,
349 + layer: torch.nn.Module,
350 + input: torch.Tensor,
351 + topk_weights: torch.Tensor,
352 + topk_ids: torch.Tensor,
353 + activation: str,
354 + global_num_experts: int = -1,
355 + ) -> torch.Tensor:
356 + output = torch.empty_like(input)
357 + layer_id = id(layer)
358 + torch.ops.vllm.cpu_fused_moe_torch(
359 + layer_id,
360 + output,
361 + input,
362 + topk_weights,
363 + topk_ids,
364 + activation,
365 + global_num_experts,
366 )
292 - return final_out
367 +
368 + return output
369 +
370 +
371 +def cpu_fused_moe_torch(
372 + layer_id: int,
373 + output: torch.Tensor,
374 + input: torch.Tensor,
375 + topk_weights: torch.Tensor,
376 + topk_ids: torch.Tensor,
377 + activation: str,
378 + global_num_experts: int = -1,
379 +) -> None:
380 + layer = _CPU_MOE_LAYER_CACHE[layer_id]()
381 +
382 + # Ref code from https://github.com/sgl-project/sglang/blob/716e682721397df103f347d22da8bd46c6016dab/python/sglang/srt/layers/moe/fused_moe_native.py#L53
383 + len_experts = global_num_experts
384 +
385 + cnts = topk_ids.new_zeros((topk_ids.shape[0], len_experts))
386 + cnts.scatter_(1, topk_ids.to(torch.int64), 1)
387 + tokens_per_expert = cnts.sum(dim=0)
388 + idxs = topk_ids.view(-1).argsort()
389 +
390 + sorted_tokens = input[idxs // topk_ids.shape[1]]
391 + tokens_per_expert = tokens_per_expert.cpu().numpy()
392 +
393 + outputs = []
394 + start_idx = 0
395 +
396 + for i, num_tokens in enumerate(tokens_per_expert):
397 + end_idx = start_idx + num_tokens
398 + if num_tokens == 0:
399 + continue
400 + tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
401 +
402 + gate_up = layer.gate_up_linear[i](tokens_for_this_expert) # type: ignore
403 + gate_up = _CPU_MOE_ACT[activation].forward_native(gate_up)
404 + expert_out = layer.down_linear[i](gate_up) # type: ignore
405 + outputs.append(expert_out)
406 + start_idx = end_idx
407 +
408 + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
409 + new_x = torch.empty_like(outs)
410 +
411 + new_x[idxs] = outs
412 + final_out = (
413 + new_x.view(*topk_ids.shape, -1)
414 + .type(topk_weights.dtype)
415 + .mul_(topk_weights.unsqueeze(dim=-1))
416 + .sum(dim=1)
417 + .type(new_x.dtype)
418 + )
419 + output.copy_(final_out)
420 +
421 +
422 +direct_register_custom_op(
423 + op_name="cpu_fused_moe_torch",
424 + op_func=cpu_fused_moe_torch,
425 + mutates_args=["output"],
426 +)
vllm/model_executor/layers/fused_moe/layer.py
+4 -2
@@ -1726,9 +1726,10 @@ class FusedMoE(CustomOp):
1726 return states
1727
1728 if self.shared_experts is None:
1729 - if current_platform.is_tpu():
1729 + if current_platform.is_tpu() or current_platform.is_cpu():
1730 # TODO: Once the OOM issue for the TPU backend is resolved, we
1731 # will switch to using the moe_forward custom op.
1732 + # Note: CPU doesn't require wrapped forward_impl.
1733 fused_output = self.forward_impl(hidden_states, router_logits)
1734 assert not isinstance(fused_output, tuple)
1735 else:
@@ -1744,9 +1745,10 @@ class FusedMoE(CustomOp):
1745 else:
1746 return reduce_output(fused_output)[..., :og_hidden_states]
1747 else:
1747 - if current_platform.is_tpu():
1748 + if current_platform.is_tpu() or current_platform.is_cpu():
1749 # TODO: Once the OOM issue for the TPU backend is resolved, we
1750 # will switch to using the moe_forward custom op.
1751 + # Note: CPU doesn't require wrapped forward_impl.
1752 shared_output, fused_output = self.forward_impl(
1753 hidden_states, router_logits
1754 )