@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0a6446005d

[CPU] Migrate unquantized MoE to the modular-kernel experts structure (#50133)

Signed-off-by: jiang1.li <jiang1.li@intel.com> Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Fadi Arafeh <fadi.arafeh@arm.com>

Li, Jiang committed Aug 3, 2026 at 14:55 UTC 0a6446005d51c9e6bfa09352f7f288ddeff17c77
20 files changed +972 -839
cmake/cpu_extension.cmake
+2 -1
@@ -461,6 +461,7 @@ set(VLLM_EXT_SRC
461 "csrc/cpu/pos_encoding.cpp"
462 "csrc/cpu/mamba_cpu.cpp"
463 "csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
464 + "csrc/cpu/cpu_fused_moe.cpp"
465 "csrc/cpu/cpu_attn.cpp"
466 "csrc/cpu/torch_bindings.cpp")
467
@@ -478,7 +479,6 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
479 "csrc/cpu/cpu_tanhf_neon.hpp"
480 ${VLLM_EXT_SRC})
481 if (ARM_BF16_FOUND)
481 - set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe.cpp" ${VLLM_EXT_SRC})
482 if (ARM_I8MM_FOUND)
483 set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe_int8.cpp" ${VLLM_EXT_SRC})
484 endif()
@@ -535,6 +535,7 @@ if (ENABLE_X86_ISA)
535
536 set(VLLM_EXT_SRC_AVX2
537 "csrc/cpu/sgl-kernels/fla.cpp"
538 + "csrc/cpu/cpu_fused_moe.cpp"
539 "csrc/cpu/utils.cpp"
540 "csrc/cpu/spec_decode_utils.cpp"
541 "csrc/cpu/cpu_attn.cpp"
csrc/cpu/cpu_fused_moe_activations.hpp
+24 -13
@@ -11,6 +11,13 @@
11 #include "cpu/cpu_arch_macros.h"
12 #include "cpu/utils.hpp"
13
14 +#if defined(DEFINE_FAST_EXP)
15 + #define DEFINE_CPU_FUSED_MOE_EXP DEFINE_FAST_EXP
16 +#else
17 + #define DEFINE_CPU_FUSED_MOE_EXP \
18 + auto fast_exp = [](const vec_op::FP32Vec16& vec) { return vec.exp(); };
19 +#endif
20 +
21 namespace cpu_fused_moe_utils {
22 enum class FusedMOEAct {
23 SiluAndMul,
@@ -39,7 +46,7 @@ void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
46 const int32_t input_stride,
47 const int32_t output_stride) {
48 using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
42 -#if !defined(__aarch64__)
49 +#if defined(__AVX512F__)
50 // For GPT-OSS interleaved gate-up weights
51 alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14,
52 16, 18, 20, 22, 24, 26, 28, 30};
@@ -50,7 +57,7 @@ void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
57 vec_op::FP32Vec16 alpha_vec(1.702);
58 vec_op::FP32Vec16 one_vec(1.0);
59
53 - DEFINE_FAST_EXP
60 + DEFINE_CPU_FUSED_MOE_EXP
61
62 for (int32_t m = 0; m < m_size; ++m) {
63 for (int32_t n = 0; n < n_size; n += 32) {
@@ -59,9 +66,19 @@ void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
66 vec_op::FP32Vec16 gate_vec(vec_op::uninit);
67 vec_op::FP32Vec16 up_vec(vec_op::uninit);
68 vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec);
62 -#else
69 +#elif defined(__AVX512F__)
70 vec_op::FP32Vec16 gate_vec(input + n, index_vec);
71 vec_op::FP32Vec16 up_vec(input + n + 1, index_vec);
72 +#else
73 + alignas(64) float gate_values[16];
74 + alignas(64) float up_values[16];
75 + const float* interleaved = input + n;
76 + for (int32_t i = 0; i < 16; ++i) {
77 + gate_values[i] = interleaved[2 * i];
78 + up_values[i] = interleaved[2 * i + 1];
79 + }
80 + vec_op::FP32Vec16 gate_vec(gate_values);
81 + vec_op::FP32Vec16 up_vec(up_values);
82 #endif
83 gate_vec = gate_vec.min(gate_up_max_vec);
84 up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec);
@@ -86,7 +103,7 @@ void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
103 float* __restrict__ up = input + dim;
104 vec_op::FP32Vec16 one_vec(1.0);
105
89 - DEFINE_FAST_EXP
106 + DEFINE_CPU_FUSED_MOE_EXP
107
108 for (int32_t m = 0; m < m_size; ++m) {
109 for (int32_t n = 0; n < dim; n += 16) {
@@ -115,21 +132,13 @@ void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
132 vec_op::FP32Vec16 one_vec(1.0);
133 vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
134 vec_op::FP32Vec16 w2_vec(0.5);
118 - alignas(64) float temp[16];
119 -
120 - DEFINE_FAST_EXP
135
136 for (int32_t m = 0; m < m_size; ++m) {
137 for (int32_t n = 0; n < dim; n += 16) {
138 vec_op::FP32Vec16 gate_vec(gate + n);
139 vec_op::FP32Vec16 up_vec(up + n);
140 auto er_input_vec = gate_vec * w1_vec;
127 -
128 - er_input_vec.save(temp);
129 - for (int32_t i = 0; i < 16; ++i) {
130 - temp[i] = std::erf(temp[i]);
131 - }
132 - vec_op::FP32Vec16 er_vec(temp);
141 + auto er_vec = er_input_vec.er();
142 auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
143 auto gated_output_fp32 = up_vec * gelu;
144 scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
@@ -201,4 +210,6 @@ FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
210 }
211 } // namespace cpu_fused_moe_utils
212
213 +#undef DEFINE_CPU_FUSED_MOE_EXP
214 +
215 #endif
csrc/cpu/cpu_types_riscv_impl.hpp
+22
@@ -678,6 +678,9 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
678 return FP32Vec16(
679 RVVI(__riscv_vfsub_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
680 }
681 + FP32Vec16 operator-() const {
682 + return FP32Vec16(RVVI(__riscv_vfneg_v_f32, LMUL_512)(reg, VEC_ELEM_NUM));
683 + }
684 FP32Vec16 operator*(const FP32Vec16& b) const {
685 return FP32Vec16(
686 RVVI(__riscv_vfmul_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
@@ -898,6 +901,25 @@ struct INT8Vec16 : public Vec<INT8Vec16> {
901 }
902 };
903
904 +// Reference implementation for vector operations missing from some backends.
905 +struct INT8Vec64 {
906 + constexpr static int VEC_ELEM_NUM = 64;
907 +
908 + explicit INT8Vec64(const int8_t* ptr) {
909 + std::memcpy(data_, ptr, sizeof(data_));
910 + }
911 +
912 + void save(int8_t* ptr) const { std::memcpy(ptr, data_, sizeof(data_)); }
913 +
914 + void save(int8_t* ptr, const int elem_num) const {
915 + TORCH_CHECK(elem_num > 0 && elem_num <= VEC_ELEM_NUM);
916 + std::memcpy(ptr, data_, elem_num);
917 + }
918 +
919 + private:
920 + int8_t data_[VEC_ELEM_NUM];
921 +};
922 +
923 // ============================================================================
924 // Type Traits & Global Helpers
925 // ============================================================================
csrc/cpu/cpu_types_scalar.hpp
+44
@@ -333,6 +333,13 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
333 return FP32Vec16(ret);
334 }
335
336 + FP32Vec16 operator-() const {
337 + f32x16_t ret;
338 + unroll_loop<int, VEC_ELEM_NUM>(
339 + [&ret, this](int i) { ret.val[i] = -reg.val[i]; });
340 + return FP32Vec16(ret);
341 + }
342 +
343 FP32Vec16 operator/(const FP32Vec16& b) const {
344 f32x16_t ret;
345 unroll_loop<int, VEC_ELEM_NUM>(
@@ -356,6 +363,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
363 return FP32Vec16(ret);
364 }
365
366 + FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
367 + return this->max(min_v).min(max_v);
368 + }
369 +
370 FP32Vec16 abs() const {
371 f32x16_t ret;
372 unroll_loop<int, VEC_ELEM_NUM>(
@@ -363,6 +374,20 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
374 return FP32Vec16(ret);
375 }
376
377 + FP32Vec16 exp() const {
378 + f32x16_t ret;
379 + unroll_loop<int, VEC_ELEM_NUM>(
380 + [&ret, this](int i) { ret.val[i] = std::exp(reg.val[i]); });
381 + return FP32Vec16(ret);
382 + }
383 +
384 + FP32Vec16 er() const {
385 + f32x16_t ret;
386 + unroll_loop<int, VEC_ELEM_NUM>(
387 + [&ret, this](int i) { ret.val[i] = std::erf(reg.val[i]); });
388 + return FP32Vec16(ret);
389 + }
390 +
391 FP32Vec16 tanh() const {
392 f32x16_t ret;
393 unroll_loop<int, VEC_ELEM_NUM>(
@@ -404,6 +429,25 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
429 void save(void* ptr) const { *reinterpret_cast<f32x16_t*>(ptr) = reg; }
430 };
431
432 +// Reference implementation for vector operations missing from some backends.
433 +struct INT8Vec64 {
434 + constexpr static int VEC_ELEM_NUM = 64;
435 +
436 + explicit INT8Vec64(const int8_t* ptr) {
437 + std::memcpy(data_, ptr, sizeof(data_));
438 + }
439 +
440 + void save(int8_t* ptr) const { std::memcpy(ptr, data_, sizeof(data_)); }
441 +
442 + void save(int8_t* ptr, const int elem_num) const {
443 + TORCH_CHECK(elem_num > 0 && elem_num <= VEC_ELEM_NUM);
444 + std::memcpy(ptr, data_, elem_num);
445 + }
446 +
447 + private:
448 + int8_t data_[VEC_ELEM_NUM];
449 +};
450 +
451 template <typename T>
452 struct VecType {
453 using vec_type = void;
csrc/cpu/cpu_types_vsx.hpp
+23
@@ -655,6 +655,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
655 vec_sub(reg.val[3], b.reg.val[3])}));
656 }
657
658 + FP32Vec16 operator-() const {
659 + return FP32Vec16(
660 + f32x4x4_t({-reg.val[0], -reg.val[1], -reg.val[2], -reg.val[3]}));
661 + }
662 +
663 FP32Vec16 operator/(const FP32Vec16& b) const {
664 return FP32Vec16(f32x4x4_t({vec_div(reg.val[0], b.reg.val[0]),
665 vec_div(reg.val[1], b.reg.val[1]),
@@ -756,6 +761,24 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
761 hi_e.reg.val[0], hi_e.reg.val[1]});
762 }
763
764 + FP32Vec16 tanh() const {
765 + FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]});
766 + FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]});
767 + auto lo_tanh = lo.tanh();
768 + auto hi_tanh = hi.tanh();
769 + return FP32Vec16(f32x4x4_t{lo_tanh.reg.val[0], lo_tanh.reg.val[1],
770 + hi_tanh.reg.val[0], hi_tanh.reg.val[1]});
771 + }
772 +
773 + FP32Vec16 er() const {
774 + FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]});
775 + FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]});
776 + auto lo_er = lo.er();
777 + auto hi_er = hi.er();
778 + return FP32Vec16(f32x4x4_t{lo_er.reg.val[0], lo_er.reg.val[1],
779 + hi_er.reg.val[0], hi_er.reg.val[1]});
780 + }
781 +
782 float reduce_max() {
783 __vector float max01 = vec_max(reg.val[0], reg.val[1]);
784 __vector float max23 = vec_max(reg.val[2], reg.val[3]);
csrc/cpu/cpu_types_vxe.hpp
+65 -1
@@ -6,8 +6,10 @@
6 #include <bit>
7 #include <cmath>
8 #include <cstdint>
9 +#include <cstring>
10 #include <limits>
11 #include <torch/all.h>
12 +
13 namespace vec_op {
14
15 struct fp8_e4m3_tag {};
@@ -727,6 +729,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
729 vec_sub(reg.val[3], b.reg.val[3])}));
730 }
731
732 + FP32Vec16 operator-() const {
733 + return FP32Vec16(f32x4x4_t({vec_neg(reg.val[0]), vec_neg(reg.val[1]),
734 + vec_neg(reg.val[2]), vec_neg(reg.val[3])}));
735 + }
736 +
737 FP32Vec16 operator/(const FP32Vec16& b) const {
738 return FP32Vec16(f32x4x4_t({vec_div(reg.val[0], b.reg.val[0]),
739 vec_div(reg.val[1], b.reg.val[1]),
@@ -734,6 +741,33 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
741 vec_div(reg.val[3], b.reg.val[3])}));
742 }
743
744 + FP32Vec16 exp() const {
745 + FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]});
746 + FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]});
747 + auto lo_exp = lo.exp();
748 + auto hi_exp = hi.exp();
749 + return FP32Vec16(f32x4x4_t{lo_exp.reg.val[0], lo_exp.reg.val[1],
750 + hi_exp.reg.val[0], hi_exp.reg.val[1]});
751 + }
752 +
753 + FP32Vec16 tanh() const {
754 + FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]});
755 + FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]});
756 + auto lo_tanh = lo.tanh();
757 + auto hi_tanh = hi.tanh();
758 + return FP32Vec16(f32x4x4_t{lo_tanh.reg.val[0], lo_tanh.reg.val[1],
759 + hi_tanh.reg.val[0], hi_tanh.reg.val[1]});
760 + }
761 +
762 + FP32Vec16 er() const {
763 + FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]});
764 + FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]});
765 + auto lo_er = lo.er();
766 + auto hi_er = hi.er();
767 + return FP32Vec16(f32x4x4_t{lo_er.reg.val[0], lo_er.reg.val[1],
768 + hi_er.reg.val[0], hi_er.reg.val[1]});
769 + }
770 +
771 float reduce_sum() const {
772 __vector float sum = vec_add(vec_add(reg.val[0], reg.val[1]),
773 vec_add(reg.val[2], reg.val[3]));
@@ -765,6 +799,17 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
799 vec_max(reg.val[3], b.reg.val[3])}));
800 }
801
802 + FP32Vec16 min(const FP32Vec16& b) const {
803 + return FP32Vec16(f32x4x4_t({vec_min(reg.val[0], b.reg.val[0]),
804 + vec_min(reg.val[1], b.reg.val[1]),
805 + vec_min(reg.val[2], b.reg.val[2]),
806 + vec_min(reg.val[3], b.reg.val[3])}));
807 + }
808 +
809 + FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
810 + return this->max(min_v).min(max_v);
811 + }
812 +
813 float reduce_max() const {
814 __vector float m = vec_max(vec_max(reg.val[0], reg.val[1]),
815 vec_max(reg.val[2], reg.val[3]));
@@ -899,6 +944,25 @@ struct INT8Vec16 : public Vec<INT8Vec16> {
944 }
945 };
946
947 +// Reference implementation for vector operations missing from some backends.
948 +struct INT8Vec64 {
949 + constexpr static int VEC_ELEM_NUM = 64;
950 +
951 + explicit INT8Vec64(const int8_t* ptr) {
952 + std::memcpy(data_, ptr, sizeof(data_));
953 + }
954 +
955 + void save(int8_t* ptr) const { std::memcpy(ptr, data_, sizeof(data_)); }
956 +
957 + void save(int8_t* ptr, const int elem_num) const {
958 + TORCH_CHECK(elem_num > 0 && elem_num <= VEC_ELEM_NUM);
959 + std::memcpy(ptr, data_, elem_num);
960 + }
961 +
962 + private:
963 + int8_t data_[VEC_ELEM_NUM];
964 +};
965 +
966 template <typename T>
967 struct VecType {
968 using vec_type = void;
@@ -1325,4 +1389,4 @@ FORCE_INLINE void prefetch(const void* addr) {
1389
1390 }; // namespace vec_op
1391
1328 -#endif
\ No newline at end of file
1392 +#endif
csrc/cpu/cpu_types_x86.hpp
+35
@@ -2,6 +2,7 @@
2 #ifndef CPU_TYPES_X86_HPP
3 #define CPU_TYPES_X86_HPP
4
5 +#include <cstring>
6 #include <immintrin.h>
7 #include <sleef.h>
8 #include <torch/all.h>
@@ -595,6 +596,8 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
596
597 FP32Vec16 tanh() const { return FP32Vec16(Sleef_tanhf16_u10(reg)); }
598
599 + FP32Vec16 er() const { return FP32Vec16(Sleef_erff16_u10(reg)); }
600 +
601 float reduce_sum() const { return _mm512_reduce_add_ps(reg); }
602
603 float reduce_max() const { return _mm512_reduce_max_ps(reg); }
@@ -798,6 +801,14 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
801 return FP32Vec16(low.tanh().reg, high.tanh().reg);
802 }
803
804 + FP32Vec16 exp() const {
805 + return FP32Vec16(Sleef_expf8_u10(reg_low), Sleef_expf8_u10(reg_high));
806 + }
807 +
808 + FP32Vec16 er() const {
809 + return FP32Vec16(Sleef_erff8_u10(reg_low), Sleef_erff8_u10(reg_high));
810 + }
811 +
812 FP32Vec16 min(const FP32Vec16& b) const {
813 return FP32Vec16(_mm256_min_ps(reg_low, b.reg_low),
814 _mm256_min_ps(reg_high, b.reg_high));
@@ -927,6 +938,30 @@ struct INT8Vec16 : public Vec<INT8Vec16> {
938 for (int i = 0; i < elem_num; ++i) ptr[i] = ar.values[i];
939 }
940 };
941 +
942 +struct INT8Vec64 : public Vec<INT8Vec64> {
943 + constexpr static int VEC_ELEM_NUM = 64;
944 +
945 + __m256i reg_low;
946 + __m256i reg_high;
947 +
948 + explicit INT8Vec64(const void* ptr)
949 + : reg_low(_mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr))),
950 + reg_high(
951 + _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr) + 1)) {}
952 +
953 + void save(void* ptr) const {
954 + _mm256_storeu_si256(reinterpret_cast<__m256i*>(ptr), reg_low);
955 + _mm256_storeu_si256(reinterpret_cast<__m256i*>(ptr) + 1, reg_high);
956 + }
957 +
958 + void save(int8_t* ptr, const int elem_num) const {
959 + TORCH_CHECK(elem_num > 0 && elem_num <= VEC_ELEM_NUM);
960 + int8_t values[VEC_ELEM_NUM];
961 + save(values);
962 + std::memcpy(ptr, values, elem_num);
963 + }
964 +};
965 #endif
966
967 template <typename T>
csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp
+1 -2
@@ -71,8 +71,7 @@ class TileGemm82 {
71 scalar_t* __restrict__ curr_m_a = curr_a;
72 vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
73 scalar_t v = *curr_m_a;
74 - load_vec_t a_reg_original(v);
75 - vec_op::FP32Vec16 a_reg(a_reg_original);
74 + vec_op::FP32Vec16 a_reg(static_cast<float>(v));
75 c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
76 c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;
77
csrc/cpu/torch_bindings.cpp
-3
@@ -614,7 +614,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
614 #endif
615
616 // fused moe
617 -#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) && !defined(__APPLE__))
617 ops.def(
618 "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
619 "-> ()");
@@ -625,8 +624,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
624 "bool skip_weighted, "
625 "str act, str isa) -> ()");
626 ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
628 -#endif // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) &&
629 - // !defined(__APPLE__))
627 #if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \
628 !defined(__APPLE__)
629 ops.def(
csrc/cpu/utils.hpp
-2
@@ -40,12 +40,10 @@ struct VecTypeTrait<c10::BFloat16> {
40 using vec_t = vec_op::BF16Vec16;
41 };
42
43 -#if !defined(__powerpc__)
43 template <>
44 struct VecTypeTrait<c10::Half> {
45 using vec_t = vec_op::FP16Vec16;
46 };
48 -#endif
47
48 struct Counter {
49 std::atomic<int64_t> counter;
docs/design/moe_kernel_features.md
+1 -1
@@ -91,7 +91,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
91 | trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],</br>[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],</br>[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],</br>[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] |
92 | hpc | standard | fp8 | G(128),T | silu | Y | Y | [`HPCExperts`][vllm.model_executor.layers.fused_moe.hpc_moe.HPCExperts] |
93 | rocm aiter moe | standard | mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai | Y | N | `rocm_aiter_fused_experts`,</br>`AiterExperts` |
94 -| cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] |
94 +| cpu_moe | standard | N/A | N/A | silu, gelu,</br>gelu_tanh,</br>swigluoai | Y | N | [`X86CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.X86CPUUnquantizedExperts],</br>[`ArmCPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.ArmCPUUnquantizedExperts],</br>[`CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.CPUUnquantizedExperts] |
95 | naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] |
96
97 !!! info "Table key"
docs/getting_started/installation/cpu.md
+1 -1
@@ -152,7 +152,7 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cpu VLLM_TARGET_DEVICE=cpu
152 - `VLLM_CPU_OMP_THREADS_BIND`: specify the CPU cores dedicated to the OpenMP threads, can be set as CPU id lists, `auto` (by default), or `nobind` (to disable binding to individual CPU cores and to inherit user-defined OpenMP variables). 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 respectively. If set to `nobind`, the number of OpenMP threads is determined by the standard `OMP_NUM_THREADS` environment variable.
153 - `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 `None`. If the value is not set and use `auto` thread binding, no CPU will be reserved for `world_size == 1`, 1 CPU per rank will be reserved for `world_size > 1`.
154 - `CPU_VISIBLE_MEMORY_NODES`: specify visible NUMA memory nodes for vLLM CPU workers, similar to ```CUDA_VISIBLE_DEVICES```. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. The variable provides more control for the auto thread-binding feature, such as masking nodes and changing nodes binding sequence.
155 -- `VLLM_CPU_SGL_KERNEL` (x86 only, 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).
155 +- `VLLM_CPU_SGL_KERNEL` (x86 only, Experimental): whether to use small-batch optimized kernels for the linear 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). MoE layers always use the grouped-gemm kernels on x86 and are unaffected by this variable.
156 - `VLLM_ZENTORCH_WEIGHT_PREPACK` (AMD Zen only): when `ZenCpuPlatform` is active, eagerly prepack linear weights into ZenDNN's blocked layout at model load time, eliminating per-inference layout conversion overhead. Default is `1` (enabled). See [AMD Zen optimizations](#amd-zen-optimizations).
157
158 ## FAQ
tests/kernels/moe/test_cpu_fused_moe.py
+161 -48
@@ -1,9 +1,12 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import sys
5 +
6 import pytest
7 import torch
8
9 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10 from tests.kernels.allclose_default import get_default_atol, get_default_rtol
11 from vllm._custom_ops import (
12 cpu_fused_moe,
@@ -12,9 +15,18 @@ from vllm._custom_ops import (
15 cpu_prepack_moe_weight_int8,
16 )
17 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
15 -from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
16 - _CPU_MOE_ACT_FN,
17 - CPUFusedMOE,
18 +from vllm.model_executor.layers.fused_moe.config import (
19 + FusedMoEConfig,
20 + FusedMoEParallelConfig,
21 + FusedMoEQuantConfig,
22 + RoutingMethodType,
23 + biased_moe_quant_config,
24 +)
25 +from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
26 + ArmCPUUnquantizedExperts,
27 + CPUUnquantizedExperts,
28 + X86CPUUnquantizedExperts,
29 + select_experts,
30 )
31 from vllm.platforms import CpuArchEnum, current_platform
32 from vllm.utils.torch_utils import set_random_seed
@@ -36,7 +48,10 @@ ACT = [
48 ]
49 USE_BIAS = [False, True]
50 ISA = ["vec"]
39 -if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
51 +if (
52 + current_platform.get_cpu_architecture() == CpuArchEnum.ARM
53 + and sys.platform != "darwin"
54 +):
55 ISA.append("neon")
56 if torch.cpu._is_amx_tile_supported():
57 ISA.append("amx")
@@ -44,6 +59,27 @@ if torch.cpu._is_amx_tile_supported():
59 DTYPE = [torch.bfloat16]
60
61
62 +def _ref_moe_activation(
63 + input: torch.Tensor,
64 + activation: MoEActivation,
65 +) -> torch.Tensor:
66 + if activation == MoEActivation.SWIGLUOAI:
67 + gate = input[..., ::2].clamp(max=7.0)
68 + up = input[..., 1::2].clamp(min=-7.0, max=7.0)
69 + glu = gate * torch.sigmoid(gate * 1.702)
70 + return (up + 1) * glu
71 +
72 + d = input.shape[-1] // 2
73 + gate, up = input[..., :d], input[..., d:]
74 + if activation == MoEActivation.SILU:
75 + return torch.nn.functional.silu(gate) * up
76 + if activation == MoEActivation.GELU:
77 + return torch.nn.functional.gelu(gate, approximate="none") * up
78 + if activation == MoEActivation.GELU_TANH:
79 + return torch.nn.functional.gelu(gate, approximate="tanh") * up
80 + raise ValueError(f"Unsupported activation: {activation}")
81 +
82 +
83 def ref_fused_moe(
84 input: torch.Tensor,
85 w13: torch.Tensor,
@@ -87,7 +123,8 @@ def ref_fused_moe(
123 tokens_for_this_expert, curr_w13, curr_w13_bias
124 )
125 # Note: to simulate the kernel implementation
90 - gate_up = _CPU_MOE_ACT_FN[activation](gate_up).to(dtype=input.dtype).float()
126 + gate_up = _ref_moe_activation(gate_up, activation)
127 + gate_up = gate_up.to(dtype=input.dtype).float()
128 expert_out = torch.nn.functional.linear(gate_up, curr_w2, curr_w2_bias)
129
130 outputs.append(expert_out)
@@ -166,7 +203,7 @@ def ref_fused_moe_int8(
203 if w13_bias is not None:
204 gate_up += w13_bias[expert_idx].float()
205
169 - intermediate = _CPU_MOE_ACT_FN[activation](gate_up).to(input.dtype)
206 + intermediate = _ref_moe_activation(gate_up, activation).to(input.dtype)
207 intermediate_int8, intermediate_scale = quantize_per_token(intermediate)
208 output = torch.matmul(
209 intermediate_int8.float(),
@@ -356,8 +393,8 @@ UNALIGNED_INTERMEDIATE_DIM = 176
393
394 class _StubMoELayer(torch.nn.Module):
395 """Minimal stand-in for the real MoE layer module, exposing just what
359 - CPUFusedMOE reads/replaces (w13_weight, w2_weight, activation, and
360 - optionally w13_bias/w2_bias)."""
396 + the unquantized CPU experts read/replace (w13_weight, w2_weight, the
397 + router configuration, and optionally w13_bias/w2_bias)."""
398
399 def __init__(
400 self,
@@ -375,12 +412,77 @@ class _StubMoELayer(torch.nn.Module):
412 self.w13_bias = torch.nn.Parameter(w13_bias, requires_grad=False)
413 if w2_bias is not None:
414 self.w2_bias = torch.nn.Parameter(w2_bias, requires_grad=False)
415 + self.use_grouped_topk = False
416 + self.renormalize = False
417 + self.scoring_func = "softmax"
418 + self.custom_routing_function = None
419 +
420 +
421 +def _make_moe_config(
422 + expert_num: int,
423 + hidden_size: int,
424 + intermediate_size: int,
425 + topk_num: int,
426 + dtype: torch.dtype,
427 + act: MoEActivation,
428 + has_bias: bool = False,
429 +) -> FusedMoEConfig:
430 + return FusedMoEConfig(
431 + num_experts=expert_num,
432 + experts_per_token=topk_num,
433 + hidden_dim=hidden_size,
434 + intermediate_size=intermediate_size,
435 + num_local_experts=expert_num,
436 + num_logical_experts=expert_num,
437 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
438 + activation=act,
439 + in_dtype=dtype,
440 + device="cpu",
441 + routing_method=RoutingMethodType.Default,
442 + has_bias=has_bias,
443 + )
444 +
445 +
446 +@pytest.mark.parametrize(
447 + ("hidden_size", "intermediate_size", "act", "expected"),
448 + [
449 + (128, 128, MoEActivation.SILU, True),
450 + # Hidden size is not 32-aligned.
451 + (112, 128, MoEActivation.SILU, False),
452 + (128, 176, MoEActivation.SILU, True),
453 + # SwiGLUOAI's interleaved gate/up layout cannot be padded.
454 + (128, 176, MoEActivation.SWIGLUOAI, False),
455 + ],
456 +)
457 +def test_cpu_vec_fused_moe_shape_support(
458 + hidden_size: int,
459 + intermediate_size: int,
460 + act: MoEActivation,
461 + expected: bool,
462 +):
463 + moe_config = _make_moe_config(
464 + expert_num=8,
465 + hidden_size=hidden_size,
466 + intermediate_size=intermediate_size,
467 + topk_num=4,
468 + dtype=torch.bfloat16,
469 + act=act,
470 + )
471 + supported, _ = CPUUnquantizedExperts.is_supported_config(
472 + CPUUnquantizedExperts,
473 + moe_config,
474 + None,
475 + None,
476 + mk.FusedMoEActivationFormat.Standard,
477 + )
478 + assert supported is expected
479
480
481 @pytest.mark.parametrize("expert_num", EXPERT_NUM)
482 @pytest.mark.parametrize("hidden_size", HIDDEN_DIM)
483 @pytest.mark.parametrize("use_bias", USE_BIAS)
484 @pytest.mark.parametrize("dtype", DTYPE)
485 +@pytest.mark.parametrize("isa", ISA)
486 @pytest.mark.parametrize(
487 "act",
488 [MoEActivation.SILU, MoEActivation.GELU, MoEActivation.GELU_TANH],
@@ -391,14 +493,11 @@ def test_cpu_fused_moe_unaligned_intermediate_size(
493 hidden_size: int,
494 use_bias: bool,
495 dtype: torch.dtype,
496 + isa: str,
497 act: MoEActivation,
498 ):
396 - """An unaligned per-partition moe_intermediate_size must still hit the
397 - grouped-gemm fast path via automatic zero-padding, with numerically
398 - correct output, instead of silently falling back to the much slower
399 - per-expert torch loop."""
400 - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
401 - pytest.skip("padding is only applied on the x86 AMX/vector kernels")
499 + """CPU kernels handle unaligned intermediate sizes by zero-padding the
500 + weights before prepacking."""
501
502 set_random_seed(0)
503 batch_size = 64
@@ -423,9 +522,15 @@ def test_cpu_fused_moe_unaligned_intermediate_size(
522 w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / (
523 0.5 * hidden_size**0.5
524 )
426 - score = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
427 - topk_weight, topk_ids = torch.topk(score, topk_num)
428 - topk_ids = topk_ids.to(torch.int32)
525 + # Route with the same helper apply() uses internally, so the reference
526 + # only differs from the kernel in how the experts are evaluated.
527 + topk_weight, topk_ids = select_experts(
528 + hidden_states=input,
529 + router_logits=router_logits,
530 + top_k=topk_num,
531 + use_grouped_topk=False,
532 + renormalize=False,
533 + )
534
535 ref_output = ref_fused_moe(
536 input, w13, w2, w13_bias, w2_bias, topk_weight, topk_ids, act
@@ -439,40 +544,48 @@ def test_cpu_fused_moe_unaligned_intermediate_size(
544 w2_bias.clone() if w2_bias is not None else None,
545 )
546
442 - cpu_moe = CPUFusedMOE(layer)
443 - assert cpu_moe.forward_method == cpu_moe.forward_grouped_gemm, (
444 - "expected the padded intermediate size to hit the grouped-gemm fast path"
547 + moe_config = _make_moe_config(
548 + expert_num, hidden_size, intermediate_size, topk_num, dtype, act, use_bias
549 + )
550 + experts_cls = {
551 + "vec": CPUUnquantizedExperts,
552 + "amx": X86CPUUnquantizedExperts,
553 + "neon": ArmCPUUnquantizedExperts,
554 + }[isa]
555 + supported, reason = experts_cls.is_supported_config(
556 + experts_cls,
557 + moe_config,
558 + None,
559 + None,
560 + mk.FusedMoEActivationFormat.Standard,
561 + )
562 + assert supported, (
563 + f"expected the padded intermediate size to be supported, got {reason}"
564 )
565
447 - output = cpu_moe.forward_method(
448 - layer, input, topk_weight, topk_ids, act, expert_num, False
566 + # Mirror UnquantizedFusedMoEMethod._setup_kernel: the quant config is
567 + # built from the layer first, then the experts shuffle the layer into the
568 + # kernel's runtime format.
569 + quant_config = (
570 + biased_moe_quant_config(layer.w13_bias, layer.w2_bias)
571 + if use_bias
572 + else FusedMoEQuantConfig.make()
573 + )
574 + experts = experts_cls(moe_config, quant_config)
575 + assert experts.isa == isa
576 + experts.process_weights_after_loading(layer)
577 +
578 + output = experts.apply(
579 + hidden_states=input,
580 + w1=layer.w13_weight,
581 + w2=layer.w2_weight,
582 + router_logits=router_logits,
583 + activation=act,
584 + global_num_experts=expert_num,
585 + expert_map=None,
586 + a1q_scale=None,
587 + apply_router_weight_on_input=False,
588 )
589
590 atol, rtol = get_default_atol(output), get_default_rtol(output)
591 torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
453 -
454 -
455 -def test_cpu_fused_moe_unaligned_intermediate_size_swigluoai(default_vllm_config):
456 - """swigluoai's interleaved gate/up layout isn't padded automatically. On
457 - AMX-capable CPUs this must raise rather than silently falling back to
458 - the (correct but much slower) per-expert torch loop; elsewhere it should
459 - fall back exactly as before."""
460 - set_random_seed(0)
461 - expert_num = 8
462 - hidden_size = 128
463 - intermediate_size = UNALIGNED_INTERMEDIATE_DIM
464 - dtype = torch.bfloat16
465 - up_dim = 2 * intermediate_size
466 -
467 - layer = _StubMoELayer(
468 - torch.randn((expert_num, up_dim, hidden_size), dtype=dtype),
469 - torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype),
470 - MoEActivation.SWIGLUOAI,
471 - )
472 -
473 - if torch.cpu._is_amx_tile_supported():
474 - with pytest.raises(RuntimeError):
475 - CPUFusedMOE(layer)
476 - else:
477 - cpu_moe = CPUFusedMOE(layer)
478 - assert cpu_moe.forward_method == cpu_moe.forward_torch
tests/kernels/moe/test_moe.py
-79
@@ -1482,85 +1482,6 @@ def test_moe_sum_pad_aware(topk: int, dtype: torch.dtype, topk_ids_dtype: torch.
1482 opcheck(torch.ops._moe_C.moe_sum, (input, actual, topk_ids, expert_map))
1483
1484
1485 -@pytest.mark.usefixtures("default_vllm_config")
1486 -@pytest.mark.parametrize("m", [1, 33])
1487 -@pytest.mark.parametrize("n,k", [(128, 128)])
1488 -@pytest.mark.parametrize("e", [8])
1489 -@pytest.mark.parametrize("topk", [2])
1490 -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
1491 -@pytest.mark.parametrize("with_bias", [False, True])
1492 -@pytest.mark.parametrize("activation", [MoEActivation.SILU])
1493 -@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only test")
1494 -def test_cpu_fused_moe_basic(
1495 - m: int,
1496 - n: int,
1497 - k: int,
1498 - e: int,
1499 - topk: int,
1500 - dtype: torch.dtype,
1501 - with_bias: bool,
1502 - activation: MoEActivation,
1503 -):
1504 - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import CPUFusedMOE
1505 -
1506 - device = "cpu"
1507 - set_random_seed(7)
1508 -
1509 - a = torch.randn((m, k), device=device, dtype=dtype) / 10
1510 - w13 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / 10
1511 - w2 = torch.randn((e, k, n), device=device, dtype=dtype) / 10
1512 - router_logits = torch.randn((m, e), device=device, dtype=dtype)
1513 -
1514 - b1 = b2 = None
1515 - if with_bias:
1516 - b1 = torch.randn((e, 2 * n), device=device, dtype=dtype) / 10
1517 - b2 = torch.randn((e, k), device=device, dtype=dtype) / 10
1518 -
1519 - ref = (
1520 - torch_moe(a, w13, w2, router_logits, topk, b1, b2)
1521 - if with_bias
1522 - else torch_moe(a, w13, w2, router_logits, topk)
1523 - )
1524 -
1525 - class _Dummy(torch.nn.Module):
1526 - def __init__(self, w13, w2, b1=None, b2=None):
1527 - super().__init__()
1528 - self.w13_weight = torch.nn.Parameter(w13, requires_grad=False)
1529 - self.w2_weight = torch.nn.Parameter(w2, requires_grad=False)
1530 - if b1 is not None:
1531 - self.w13_bias = torch.nn.Parameter(b1, requires_grad=False)
1532 - if b2 is not None:
1533 - self.w2_bias = torch.nn.Parameter(b2, requires_grad=False)
1534 -
1535 - layer = _Dummy(w13, w2, b1, b2).to(dtype)
1536 - fused = CPUFusedMOE(layer)
1537 - out = fused(
1538 - layer=layer,
1539 - x=a,
1540 - use_grouped_topk=False,
1541 - top_k=topk,
1542 - router_logits=router_logits,
1543 - renormalize=False,
1544 - global_num_experts=e,
1545 - expert_map=None,
1546 - custom_routing_function=None,
1547 - scoring_func="softmax",
1548 - routed_scaling_factor=1.0,
1549 - e_score_correction_bias=None,
1550 - apply_router_weight_on_input=False,
1551 - activation=activation,
1552 - )
1553 -
1554 - # Tolerances: fp32 tight; bf16 looser (esp. with bias)
1555 - if dtype == torch.float32:
1556 - atol = 1e-3
1557 - elif with_bias:
1558 - atol = 8e-2
1559 - else:
1560 - atol = 5e-2
1561 - torch.testing.assert_close(out, ref, atol=atol, rtol=0)
1562 -
1563 -
1485 def _batched_fused_marlin_moe_cases() -> list[Any]:
1486 cases = [
1487 pytest.param(
tests/kernels/moe/test_unquantized_backend_selection.py
+114 -3
@@ -3,14 +3,16 @@
3 from unittest.mock import patch
4
5 import pytest
6 +import torch
7
8 from tests.kernels.moe.utils import make_dummy_moe_config
9 from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
10 from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
11 UnquantizedMoeBackend,
12 + backend_to_kernel_cls,
13 select_unquantized_moe_backend,
14 )
13 -from vllm.platforms import current_platform
15 +from vllm.platforms import CpuArchEnum, current_platform
16
17 skipif_not_cuda_rocm = pytest.mark.skipif(
18 not (current_platform.is_cuda() or current_platform.is_rocm()),
@@ -18,6 +20,109 @@ skipif_not_cuda_rocm = pytest.mark.skipif(
20 )
21
22
23 +@pytest.mark.parametrize(
24 + ("amx_supported", "in_dtype", "expect_amx_kernel"),
25 + [
26 + (True, torch.bfloat16, True),
27 + (False, torch.bfloat16, False),
28 + (True, torch.float16, False),
29 + ],
30 +)
31 +def test_x86_cpu_unquantized_kernel_selection(
32 + amx_supported: bool,
33 + in_dtype: torch.dtype,
34 + expect_amx_kernel: bool,
35 +):
36 + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
37 + CPUUnquantizedExperts,
38 + X86CPUUnquantizedExperts,
39 + )
40 +
41 + with (
42 + patch.object(current_platform, "is_cpu", return_value=True),
43 + patch.object(
44 + current_platform,
45 + "get_cpu_architecture",
46 + return_value=CpuArchEnum.X86,
47 + ),
48 + patch("torch.cpu._is_amx_tile_supported", return_value=amx_supported),
49 + ):
50 + moe_config = make_dummy_moe_config(
51 + hidden_dim=128,
52 + intermediate_size=128,
53 + in_dtype=in_dtype,
54 + )
55 + kernel_cls = next(
56 + cls
57 + for cls in backend_to_kernel_cls(UnquantizedMoeBackend.CPU)
58 + if cls.is_supported_config(
59 + cls,
60 + moe_config,
61 + None,
62 + None,
63 + CPUUnquantizedExperts.activation_format(),
64 + )[0]
65 + )
66 +
67 + expected_kernel_cls = (
68 + X86CPUUnquantizedExperts if expect_amx_kernel else CPUUnquantizedExperts
69 + )
70 + assert kernel_cls is expected_kernel_cls
71 +
72 +
73 +@pytest.mark.parametrize(
74 + ("platform", "in_dtype", "expect_arm_kernel"),
75 + [
76 + ("linux", torch.bfloat16, True),
77 + ("linux", torch.float16, False),
78 + ("darwin", torch.bfloat16, False),
79 + ],
80 +)
81 +def test_arm_cpu_unquantized_kernel_selection(
82 + platform: str,
83 + in_dtype: torch.dtype,
84 + expect_arm_kernel: bool,
85 +):
86 + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
87 + ArmCPUUnquantizedExperts,
88 + CPUUnquantizedExperts,
89 + )
90 +
91 + with (
92 + patch.object(current_platform, "is_cpu", return_value=True),
93 + patch.object(
94 + current_platform,
95 + "get_cpu_architecture",
96 + return_value=CpuArchEnum.ARM,
97 + ),
98 + patch(
99 + "vllm.model_executor.layers.fused_moe.experts.cpu_moe.sys.platform",
100 + platform,
101 + ),
102 + ):
103 + moe_config = make_dummy_moe_config(
104 + hidden_dim=128,
105 + intermediate_size=128,
106 + in_dtype=in_dtype,
107 + )
108 + kernel_cls = next(
109 + cls
110 + for cls in backend_to_kernel_cls(UnquantizedMoeBackend.CPU)
111 + if cls.is_supported_config(
112 + cls,
113 + moe_config,
114 + None,
115 + None,
116 + CPUUnquantizedExperts.activation_format(),
117 + )[0]
118 + )
119 +
120 + expected_kernel_cls = (
121 + ArmCPUUnquantizedExperts if expect_arm_kernel else CPUUnquantizedExperts
122 + )
123 + assert kernel_cls is expected_kernel_cls
124 +
125 +
126 @pytest.mark.parametrize(
127 "platform_method,expected_backend",
128 [
@@ -69,14 +174,20 @@ def test_select_default_backend_by_platform(
174 patch.object(current_platform, "is_out_of_tree", return_value=False),
175 patch.object(current_platform, platform_method, return_value=True),
176 ):
72 - moe_config = make_dummy_moe_config()
177 + # CPU's grouped-gemm kernels require hidden/intermediate sizes
178 + # aligned to 32; the size-1 defaults only work for backends that
179 + # don't check shapes at selection time.
180 + moe_config = (
181 + make_dummy_moe_config(hidden_dim=128, intermediate_size=128)
182 + if expected_backend == UnquantizedMoeBackend.CPU
183 + else make_dummy_moe_config()
184 + )
185 selected_backend, expert_cls = select_unquantized_moe_backend(
186 moe_config=moe_config
187 )
188
189 assert selected_backend == expected_backend
190 if expected_backend in [
79 - UnquantizedMoeBackend.CPU,
191 UnquantizedMoeBackend.OOT,
192 UnquantizedMoeBackend.TPU,
193 ]:
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py deleted
-577
@@ -1,577 +0,0 @@
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 (
11 - CPUQuantMethod,
12 - cpu_fused_moe,
13 - cpu_prepack_moe_weight,
14 - fused_experts_cpu,
15 -)
16 -from vllm.model_executor.layers.activation import SiluAndMul
17 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation
18 -from vllm.model_executor.layers.quantization.utils.layer_utils import replace_parameter
19 -from vllm.platforms import CpuArchEnum, current_platform
20 -from vllm.utils.torch_utils import direct_register_custom_op
21 -
22 -_CPU_MOE_LAYER_CACHE = {}
23 -# The CPU grouped-gemm MoE kernels (AMX and vector) tile the expert
24 -# intermediate ("N") dimension in blocks of this size and have no tail/
25 -# remainder handling, so a shard is only eligible for the fast path when
26 -# its per-partition intermediate size is a multiple of it.
27 -_MOE_GROUPED_GEMM_N_TILE = 32
28 -
29 -
30 -def _swigluoai_forward_native(
31 - x: torch.Tensor,
32 - alpha: float = 1.702,
33 - limit: float = 7.0,
34 -) -> torch.Tensor:
35 - """PyTorch-native implementation of SwigluOAIAndMul.forward_native.
36 -
37 - Standalone function to avoid instantiating SwigluOAIAndMul (a CustomOp)
38 - which would trigger get_current_vllm_config() before config is set.
39 - """
40 - gate, up = x[..., ::2], x[..., 1::2]
41 - gate = gate.clamp(min=None, max=limit)
42 - up = up.clamp(min=-limit, max=limit)
43 - glu = gate * torch.sigmoid(gate * alpha)
44 - gated_output = (up + 1) * glu
45 - return gated_output
46 -
47 -
48 -def _gelu_and_mul(
49 - x: torch.Tensor,
50 -) -> torch.Tensor:
51 - d = x.shape[-1] // 2
52 - return F.gelu(x[..., :d], approximate="none") * x[..., d:]
53 -
54 -
55 -# Map activation names to their native forward functions.
56 -# Uses static methods or standalone functions to avoid instantiating CustomOp
57 -# classes, which would call get_current_vllm_config() before config is set.
58 -_CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = {
59 - MoEActivation.SILU: SiluAndMul.forward_native,
60 - MoEActivation.SWIGLUOAI: _swigluoai_forward_native,
61 - MoEActivation.GELU: _gelu_and_mul,
62 - MoEActivation.GELU_TANH: (
63 - lambda x: F.gelu(x[..., : x.shape[-1] // 2], approximate="tanh")
64 - * x[..., x.shape[-1] // 2 :]
65 - ),
66 -}
67 -
68 -
69 -def grouped_topk(
70 - hidden_states: torch.Tensor,
71 - gating_output: torch.Tensor,
72 - topk: int,
73 - renormalize: bool,
74 - num_expert_group: int = 0,
75 - topk_group: int = 0,
76 - scoring_func: str = "softmax",
77 - routed_scaling_factor: float = 1.0,
78 - e_score_correction_bias: torch.Tensor | None = None,
79 -) -> tuple[torch.Tensor, torch.Tensor]:
80 - assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
81 -
82 - gating_output = gating_output.float()
83 - if scoring_func == "softmax":
84 - scores = torch.softmax(gating_output, dim=-1)
85 - elif scoring_func == "sigmoid":
86 - scores = gating_output.sigmoid()
87 - else:
88 - raise ValueError(f"Unsupported scoring function: {scoring_func}")
89 -
90 - num_token = scores.shape[0]
91 - if e_score_correction_bias is not None:
92 - original_scores = scores
93 - scores = scores + e_score_correction_bias.unsqueeze(0)
94 - group_scores = (
95 - scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
96 - )
97 - else:
98 - group_scores = (
99 - scores.view(num_token, num_expert_group, -1).max(dim=-1).values
100 - ) # [n, n_group]
101 - group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
102 - 1
103 - ] # [n, top_k_group]
104 - group_mask = torch.zeros_like(group_scores) # [n, n_group]
105 - group_mask.scatter_(1, group_idx, 1) # [n, n_group]
106 - score_mask = (
107 - group_mask.unsqueeze(-1)
108 - .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
109 - .reshape(num_token, -1)
110 - ) # [n, e]
111 - tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) # [n, e]
112 -
113 - if e_score_correction_bias is not None:
114 - topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1]
115 - topk_weights = original_scores.gather(1, topk_ids)
116 - else:
117 - topk_weights, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)
118 -
119 - if renormalize:
120 - topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
121 -
122 - if routed_scaling_factor != 1.0:
123 - topk_weights = topk_weights * routed_scaling_factor
124 - return topk_weights, topk_ids.to(torch.int32)
125 -
126 -
127 -def select_experts(
128 - hidden_states: torch.Tensor,
129 - router_logits: torch.Tensor,
130 - top_k: int,
131 - use_grouped_topk: bool,
132 - renormalize: bool,
133 - topk_group: int | None = None,
134 - num_expert_group: int | None = None,
135 - custom_routing_function: Callable | None = None,
136 - scoring_func: str = "softmax",
137 - routed_scaling_factor: float = 1.0,
138 - e_score_correction_bias: torch.Tensor | None = None,
139 -) -> tuple[torch.Tensor, torch.Tensor]:
140 - if use_grouped_topk:
141 - assert topk_group is not None
142 - assert num_expert_group is not None
143 - return grouped_topk(
144 - hidden_states=hidden_states,
145 - gating_output=router_logits,
146 - topk=top_k,
147 - renormalize=renormalize,
148 - num_expert_group=num_expert_group,
149 - topk_group=topk_group,
150 - scoring_func=scoring_func,
151 - routed_scaling_factor=routed_scaling_factor,
152 - e_score_correction_bias=e_score_correction_bias,
153 - )
154 - elif custom_routing_function is None:
155 - assert scoring_func == "softmax"
156 - topk_logit_vals, topk_idx = torch.topk(
157 - router_logits, k=top_k, dim=-1, sorted=False
158 - )
159 - if renormalize:
160 - topk_vals = torch.softmax(topk_logit_vals, dim=-1)
161 - else:
162 - logZ = torch.logsumexp(router_logits, dim=-1, keepdim=True)
163 - topk_vals = (topk_logit_vals - logZ).exp()
164 - return topk_vals.to(torch.float32), topk_idx.to(torch.int32)
165 - else:
166 - return custom_routing_function(
167 - hidden_states=hidden_states,
168 - gating_output=router_logits,
169 - topk=top_k,
170 - renormalize=renormalize,
171 - )
172 -
173 -
174 -class SGLFusedMOE:
175 - def __init__(self, layer: torch.nn.Module) -> None:
176 - pass
177 -
178 - def __call__(
179 - self,
180 - layer: torch.nn.Module,
181 - x: torch.Tensor,
182 - use_grouped_topk: bool,
183 - top_k: int,
184 - router_logits: torch.Tensor,
185 - renormalize: bool,
186 - topk_group: int | None = None,
187 - num_expert_group: int | None = None,
188 - global_num_experts: int = -1,
189 - expert_map: torch.Tensor | None = None,
190 - custom_routing_function: Callable | None = None,
191 - scoring_func: str = "softmax",
192 - routed_scaling_factor: float = 1.0,
193 - e_score_correction_bias: torch.Tensor | None = None,
194 - apply_router_weight_on_input: bool = False,
195 - activation: MoEActivation = MoEActivation.SILU,
196 - ) -> torch.Tensor:
197 - assert activation == MoEActivation.SILU, f"{activation} is not supported."
198 - assert not apply_router_weight_on_input
199 - topk_weights, topk_ids = select_experts(
200 - hidden_states=x,
201 - router_logits=router_logits,
202 - use_grouped_topk=use_grouped_topk,
203 - top_k=top_k,
204 - renormalize=renormalize,
205 - topk_group=topk_group,
206 - num_expert_group=num_expert_group,
207 - custom_routing_function=custom_routing_function,
208 - scoring_func=scoring_func,
209 - routed_scaling_factor=routed_scaling_factor,
210 - e_score_correction_bias=e_score_correction_bias,
211 - )
212 -
213 - return fused_experts_cpu(
214 - x,
215 - layer.w13_weight,
216 - layer.w2_weight,
217 - topk_weights,
218 - topk_ids,
219 - False, # inplace
220 - CPUQuantMethod.UNQUANT, # moe_comp_method
221 - None, # w1_scale
222 - None, # w2_scale
223 - None, # w1_zero
224 - None, # w2_zero
225 - None, # block_size
226 - None, # w1_bias
227 - None, # w2_bias
228 - None, # alpha
229 - None, # limit
230 - True, # is_vnni
231 - )
232 -
233 -
234 -class CPUFusedMOE:
235 - """CPU-based fused MoE implementation."""
236 -
237 - def __init__(self, layer: torch.nn.Module) -> None:
238 - self._pad_moe_intermediate_for_grouped_gemm(layer)
239 - use_grouped_gemm, isa = self.check_grouped_gemm(layer)
240 - self.isa = isa
241 - if use_grouped_gemm:
242 - self.forward_method = self.forward_grouped_gemm
243 - self.init_moe_grouped_gemm(layer=layer)
244 - else:
245 - self.forward_method = self.forward_torch
246 - self.init_moe_torch(layer=layer)
247 -
248 - def __call__(
249 - self,
250 - layer: torch.nn.Module,
251 - x: torch.Tensor,
252 - use_grouped_topk: bool,
253 - top_k: int,
254 - router_logits: torch.Tensor,
255 - renormalize: bool,
256 - topk_group: int | None = None,
257 - num_expert_group: int | None = None,
258 - global_num_experts: int = -1,
259 - expert_map: torch.Tensor | None = None,
260 - custom_routing_function: Callable | None = None,
261 - scoring_func: str = "softmax",
262 - routed_scaling_factor: float = 1.0,
263 - e_score_correction_bias: torch.Tensor | None = None,
264 - apply_router_weight_on_input: bool = False,
265 - activation: MoEActivation = MoEActivation.SILU,
266 - ) -> torch.Tensor:
267 - assert activation in _CPU_MOE_ACT_FN, f"{activation} is not supported."
268 -
269 - topk_weights, topk_ids = select_experts(
270 - hidden_states=x,
271 - router_logits=router_logits,
272 - use_grouped_topk=use_grouped_topk,
273 - top_k=top_k,
274 - renormalize=renormalize,
275 - topk_group=topk_group,
276 - num_expert_group=num_expert_group,
277 - custom_routing_function=custom_routing_function,
278 - scoring_func=scoring_func,
279 - routed_scaling_factor=routed_scaling_factor,
280 - e_score_correction_bias=e_score_correction_bias,
281 - )
282 -
283 - return self.forward_method(
284 - layer,
285 - x,
286 - topk_weights,
287 - topk_ids,
288 - activation,
289 - global_num_experts,
290 - apply_router_weight_on_input,
291 - )
292 -
293 - def _pad_moe_intermediate_for_grouped_gemm(self, layer: torch.nn.Module) -> None:
294 - """Zero-pad the per-partition MoE intermediate dim up to a multiple
295 - of _MOE_GROUPED_GEMM_N_TILE, so the AMX/vector grouped-gemm kernels
296 - can be used even when TP sharding (moe_intermediate_size // tp_size)
297 - lands on an unaligned value (e.g. moe_intermediate_size=704 at tp=4
298 - -> 176). Only applies to the x86 (AMX/vec) kernels and half-split
299 - gate/up activations; interleaved layouts (swigluoai) are left
300 - untouched.
301 - """
302 - if not hasattr(torch.ops._C, "prepack_moe_weight"):
303 - return
304 - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
305 - return
306 -
307 - intermediate_size = layer.w2_weight.size(2)
308 - remainder = intermediate_size % _MOE_GROUPED_GEMM_N_TILE
309 - if remainder == 0:
310 - return
311 - if layer.activation == MoEActivation.SWIGLUOAI:
312 - return
313 -
314 - pad = _MOE_GROUPED_GEMM_N_TILE - remainder
315 - padded_size = intermediate_size + pad
316 - num_experts, _, hidden_size = layer.w13_weight.shape
317 -
318 - new_w13 = layer.w13_weight.new_zeros(num_experts, 2 * padded_size, hidden_size)
319 - new_w13[:, :intermediate_size] = layer.w13_weight[:, :intermediate_size]
320 - new_w13[:, padded_size : padded_size + intermediate_size] = layer.w13_weight[
321 - :, intermediate_size:
322 - ]
323 - replace_parameter(layer, "w13_weight", new_w13)
324 -
325 - new_w2 = layer.w2_weight.new_zeros(num_experts, hidden_size, padded_size)
326 - new_w2[:, :, :intermediate_size] = layer.w2_weight
327 - replace_parameter(layer, "w2_weight", new_w2)
328 -
329 - if hasattr(layer, "w13_bias"):
330 - new_bias = layer.w13_bias.new_zeros(num_experts, 2 * padded_size)
331 - new_bias[:, :intermediate_size] = layer.w13_bias[:, :intermediate_size]
332 - new_bias[:, padded_size : padded_size + intermediate_size] = layer.w13_bias[
333 - :, intermediate_size:
334 - ]
335 - replace_parameter(layer, "w13_bias", new_bias)
336 -
337 - def _grouped_gemm_alignment_error(self, layer: torch.nn.Module) -> str:
338 - # w2's input size is the per-partition MoE intermediate size (the
339 - # dimension TP-sharding splits), and it's what most commonly breaks
340 - # alignment, e.g. moe_intermediate_size=704 at tp=4 gives
341 - # 704 // 4 == 176, which isn't a multiple of 32.
342 - intermediate_size_per_partition = layer.w2_weight.size(2)
343 - return (
344 - "CPU fused-MoE AMX grouped-gemm kernel cannot be used for a "
345 - f"layer with w13 shape {tuple(layer.w13_weight.shape)} / w2 "
346 - f"shape {tuple(layer.w2_weight.shape)}: the per-partition MoE "
347 - f"intermediate size ({intermediate_size_per_partition}) is not "
348 - f"a multiple of {_MOE_GROUPED_GEMM_N_TILE}, and automatic "
349 - "zero-padding could not resolve this (typically because the "
350 - "activation uses an interleaved gate/up layout, e.g. "
351 - "swigluoai). vLLM refuses to silently fall back to the much "
352 - "slower per-expert torch loop on AMX-capable CPUs; consider a "
353 - "different --tensor-parallel-size."
354 - )
355 -
356 - def check_grouped_gemm(
357 - self,
358 - layer: torch.nn.Module,
359 - ) -> tuple[bool, str]:
360 - if not hasattr(torch.ops._C, "prepack_moe_weight"):
361 - return False, "none"
362 -
363 - dtype = layer.w13_weight.dtype
364 - w13_input_size = layer.w13_weight.size(2)
365 - w13_output_size = layer.w13_weight.size(1)
366 - w2_input_size = layer.w2_weight.size(2)
367 - w2_output_size = layer.w2_weight.size(1)
368 -
369 - supports_amx = torch.cpu._is_amx_tile_supported()
370 - if supports_amx:
371 - if (
372 - dtype == torch.bfloat16
373 - and w13_output_size % 32 == 0
374 - and w2_output_size % 32 == 0
375 - and w13_input_size % 32 == 0
376 - and w2_input_size % 32 == 0
377 - ):
378 - return True, "amx"
379 - raise RuntimeError(self._grouped_gemm_alignment_error(layer))
380 -
381 - if not (w13_output_size % 32 == 0 and w2_output_size % 32 == 0):
382 - return False, "none"
383 -
384 - if (
385 - layer.activation == MoEActivation.SWIGLUOAI
386 - and w2_input_size % _MOE_GROUPED_GEMM_N_TILE != 0
387 - ):
388 - return False, "none"
389 -
390 - supports_neon = current_platform.get_cpu_architecture() == CpuArchEnum.ARM
391 - if supports_neon:
392 - if (
393 - dtype == torch.bfloat16
394 - and w13_input_size % 4 == 0
395 - and w2_input_size % 4 == 0
396 - ):
397 - return True, "neon"
398 - return False, "none"
399 -
400 - return True, "vec"
401 -
402 - def init_moe_grouped_gemm(
403 - self,
404 - layer: torch.nn.Module,
405 - ) -> None:
406 - new_w13 = cpu_prepack_moe_weight(layer.w13_weight, self.isa)
407 - replace_parameter(layer, "w13_weight", new_w13)
408 - new_w2 = cpu_prepack_moe_weight(layer.w2_weight, self.isa)
409 - replace_parameter(layer, "w2_weight", new_w2)
410 -
411 - def init_moe_torch(
412 - self,
413 - layer: torch.nn.Module,
414 - ) -> None:
415 - use_onednn_mm = ops._supports_onednn and ops.is_onednn_acl_supported()
416 - num_experts = layer.w13_weight.size(0)
417 - has_w13_bias = hasattr(layer, "w13_bias")
418 - has_w2_bias = hasattr(layer, "w2_bias")
419 -
420 - layer.gate_up_linear = []
421 - layer.down_linear = []
422 -
423 - for i in range(num_experts):
424 - layer_w13_weight = layer.w13_weight[i]
425 - layer_w13_bias = layer.w13_bias[i] if has_w13_bias else None
426 - layer_w2_weight = layer.w2_weight[i]
427 - layer_w2_bias = layer.w2_bias[i] if has_w2_bias else None
428 - if use_onednn_mm:
429 - gate_up_handle = ops.create_onednn_mm(layer_w13_weight.t(), 32)
430 - layer.gate_up_linear.append(
431 - lambda x, handle=gate_up_handle, bias=layer_w13_bias: ops.onednn_mm(
432 - handle, x, bias
433 - )
434 - )
435 - down_handle = ops.create_onednn_mm(layer_w2_weight.t(), 32)
436 - layer.down_linear.append(
437 - lambda x, handle=down_handle, bias=layer_w2_bias: ops.onednn_mm(
438 - handle, x, bias
439 - )
440 - )
441 - else:
442 - layer.gate_up_linear.append(
443 - lambda x, w=layer_w13_weight, b=layer_w13_bias: F.linear(x, w, b)
444 - )
445 - layer.down_linear.append(
446 - lambda x, w=layer_w2_weight, b=layer_w2_bias: F.linear(x, w, b)
447 - )
448 -
449 - if use_onednn_mm: # remove weight
450 - layer.w13_weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
451 - layer.w2_weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
452 -
453 - _CPU_MOE_LAYER_CACHE[id(layer)] = weakref.ref(layer)
454 -
455 - def forward_grouped_gemm(
456 - self,
457 - layer: torch.nn.Module,
458 - input: torch.Tensor,
459 - topk_weights: torch.Tensor,
460 - topk_ids: torch.Tensor,
461 - activation: MoEActivation,
462 - global_num_experts: int = -1,
463 - skip_weighted: bool = False,
464 - ) -> torch.Tensor:
465 - if skip_weighted:
466 - assert topk_ids.size(1) == 1, (
467 - "apply_router_weight_on_input is only implemented for topk=1"
468 - )
469 - input.mul_(topk_weights.to(input.dtype))
470 -
471 - output = cpu_fused_moe(
472 - input,
473 - layer.w13_weight,
474 - layer.w2_weight,
475 - getattr(layer, "w13_bias", None),
476 - getattr(layer, "w2_bias", None),
477 - topk_weights,
478 - topk_ids,
479 - activation.value,
480 - self.isa,
481 - skip_weighted,
482 - )
483 - return output
484 -
485 - def forward_torch(
486 - self,
487 - layer: torch.nn.Module,
488 - input: torch.Tensor,
489 - topk_weights: torch.Tensor,
490 - topk_ids: torch.Tensor,
491 - activation: MoEActivation,
492 - global_num_experts: int = -1,
493 - skip_weighted: bool = False,
494 - ) -> torch.Tensor:
495 - if skip_weighted:
496 - assert topk_ids.size(1) == 1, (
497 - "apply_router_weight_on_input is only implemented for topk=1"
498 - )
499 - input.mul_(topk_weights.to(input.dtype))
500 -
501 - output = torch.empty_like(input)
502 - layer_id = id(layer)
503 - torch.ops.vllm.cpu_fused_moe_torch(
504 - layer_id,
505 - output,
506 - input,
507 - topk_weights,
508 - topk_ids,
509 - activation.value,
510 - global_num_experts,
511 - skip_weighted,
512 - )
513 -
514 - return output
515 -
516 -
517 -def cpu_fused_moe_torch(
518 - layer_id: int,
519 - output: torch.Tensor,
520 - input: torch.Tensor,
521 - topk_weights: torch.Tensor,
522 - topk_ids: torch.Tensor,
523 - activation: str,
524 - global_num_experts: int = -1,
525 - skip_weighted: bool = False,
526 -) -> None:
527 - act = MoEActivation.from_str(activation)
528 - layer = _CPU_MOE_LAYER_CACHE[layer_id]()
529 -
530 - # Ref code from https://github.com/sgl-project/sglang/blob/716e682721397df103f347d22da8bd46c6016dab/python/sglang/srt/layers/moe/fused_moe_native.py#L53
531 - len_experts = global_num_experts
532 -
533 - cnts = topk_ids.new_zeros((topk_ids.shape[0], len_experts))
534 - cnts.scatter_(1, topk_ids.to(torch.int64), 1)
535 - tokens_per_expert = cnts.sum(dim=0)
536 - idxs = topk_ids.view(-1).argsort()
537 -
538 - sorted_tokens = input[idxs // topk_ids.shape[1]]
539 - tokens_per_expert = tokens_per_expert.cpu().numpy()
540 -
541 - outputs = []
542 - start_idx = 0
543 -
544 - for i, num_tokens in enumerate(tokens_per_expert):
545 - end_idx = start_idx + num_tokens
546 - if num_tokens == 0:
547 - continue
548 - tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
549 -
550 - gate_up = layer.gate_up_linear[i](tokens_for_this_expert) # type: ignore
551 - gate_up = _CPU_MOE_ACT_FN[act](gate_up)
552 - expert_out = layer.down_linear[i](gate_up) # type: ignore
553 - outputs.append(expert_out)
554 - start_idx = end_idx
555 -
556 - outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
557 - new_x = torch.empty_like(outs)
558 -
559 - new_x[idxs] = outs
560 - if skip_weighted:
561 - final_out = new_x
562 - else:
563 - final_out = (
564 - new_x.view(*topk_ids.shape, -1)
565 - .type(topk_weights.dtype)
566 - .mul_(topk_weights.unsqueeze(dim=-1))
567 - .sum(dim=1)
568 - .type(new_x.dtype)
569 - )
570 - output.copy_(final_out)
571 -
572 -
573 -direct_register_custom_op(
574 - op_name="cpu_fused_moe_torch",
575 - op_func=cpu_fused_moe_torch,
576 - mutates_args=["output"],
577 -)
vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py
+1 -1
@@ -218,7 +218,7 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
218 Returns:
219 Output tensor after MoE computation
220 """
221 - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
221 + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
222 select_experts,
223 )
224
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+440 -24
@@ -1,6 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""CPU quantized fused MoE experts."""
3 +"""CPU fused MoE experts."""
4 +
5 +import math
6 +import sys
7 +from collections.abc import Callable
8 +from typing import cast
9
10 import torch
11
@@ -9,7 +14,9 @@ from vllm._custom_ops import (
14 CPUQuantAlgo,
15 CPUQuantMethod,
16 convert_weight_packed_scale_zp,
17 + cpu_fused_moe,
18 cpu_fused_moe_int8,
19 + cpu_prepack_moe_weight,
20 cpu_prepack_moe_weight_int8,
21 fused_experts_cpu,
22 )
@@ -29,7 +36,439 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
36 kInt8StaticChannelSym,
37 kMxfp4Static,
38 )
39 +from vllm.model_executor.utils import replace_parameter
40 from vllm.platforms import CpuArchEnum, current_platform
41 +from vllm.utils.math_utils import round_up
42 +
43 +# ===========================================================================
44 +# Routing
45 +# ===========================================================================
46 +
47 +
48 +def grouped_topk(
49 + hidden_states: torch.Tensor,
50 + gating_output: torch.Tensor,
51 + topk: int,
52 + renormalize: bool,
53 + num_expert_group: int = 0,
54 + topk_group: int = 0,
55 + scoring_func: str = "softmax",
56 + routed_scaling_factor: float = 1.0,
57 + e_score_correction_bias: torch.Tensor | None = None,
58 +) -> tuple[torch.Tensor, torch.Tensor]:
59 + assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
60 +
61 + gating_output = gating_output.float()
62 + if scoring_func == "softmax":
63 + scores = torch.softmax(gating_output, dim=-1)
64 + elif scoring_func == "sigmoid":
65 + scores = gating_output.sigmoid()
66 + else:
67 + raise ValueError(f"Unsupported scoring function: {scoring_func}")
68 +
69 + num_token = scores.shape[0]
70 + if e_score_correction_bias is not None:
71 + original_scores = scores
72 + scores = scores + e_score_correction_bias.unsqueeze(0)
73 + group_scores = (
74 + scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
75 + )
76 + else:
77 + group_scores = (
78 + scores.view(num_token, num_expert_group, -1).max(dim=-1).values
79 + ) # [n, n_group]
80 + group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
81 + 1
82 + ] # [n, top_k_group]
83 + group_mask = torch.zeros_like(group_scores) # [n, n_group]
84 + group_mask.scatter_(1, group_idx, 1) # [n, n_group]
85 + score_mask = (
86 + group_mask.unsqueeze(-1)
87 + .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
88 + .reshape(num_token, -1)
89 + ) # [n, e]
90 + tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) # [n, e]
91 +
92 + if e_score_correction_bias is not None:
93 + topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1]
94 + topk_weights = original_scores.gather(1, topk_ids)
95 + else:
96 + topk_weights, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)
97 +
98 + if renormalize:
99 + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
100 +
101 + if routed_scaling_factor != 1.0:
102 + topk_weights = topk_weights * routed_scaling_factor
103 + return topk_weights, topk_ids.to(torch.int32)
104 +
105 +
106 +def select_experts(
107 + hidden_states: torch.Tensor,
108 + router_logits: torch.Tensor,
109 + top_k: int,
110 + use_grouped_topk: bool,
111 + renormalize: bool,
112 + topk_group: int | None = None,
113 + num_expert_group: int | None = None,
114 + custom_routing_function: Callable | None = None,
115 + scoring_func: str = "softmax",
116 + routed_scaling_factor: float = 1.0,
117 + e_score_correction_bias: torch.Tensor | None = None,
118 +) -> tuple[torch.Tensor, torch.Tensor]:
119 + if use_grouped_topk:
120 + assert topk_group is not None
121 + assert num_expert_group is not None
122 + return grouped_topk(
123 + hidden_states=hidden_states,
124 + gating_output=router_logits,
125 + topk=top_k,
126 + renormalize=renormalize,
127 + num_expert_group=num_expert_group,
128 + topk_group=topk_group,
129 + scoring_func=scoring_func,
130 + routed_scaling_factor=routed_scaling_factor,
131 + e_score_correction_bias=e_score_correction_bias,
132 + )
133 + elif custom_routing_function is None:
134 + assert scoring_func == "softmax"
135 + topk_logit_vals, topk_idx = torch.topk(
136 + router_logits, k=top_k, dim=-1, sorted=False
137 + )
138 + if renormalize:
139 + topk_vals = torch.softmax(topk_logit_vals, dim=-1)
140 + else:
141 + logZ = torch.logsumexp(router_logits, dim=-1, keepdim=True)
142 + topk_vals = (topk_logit_vals - logZ).exp()
143 + return topk_vals.to(torch.float32), topk_idx.to(torch.int32)
144 + else:
145 + topk_weights, topk_ids = custom_routing_function(
146 + hidden_states=hidden_states,
147 + gating_output=router_logits,
148 + topk=top_k,
149 + renormalize=renormalize,
150 + )
151 + # cpu_fused_moe reads routing tensors as contiguous float32/int32
152 + # buffers and does not account for tensor strides.
153 + topk_weights = topk_weights.to(torch.float32).contiguous()
154 + topk_ids = topk_ids.to(torch.int32).contiguous()
155 + return topk_weights, topk_ids
156 +
157 +
158 +# ===========================================================================
159 +# Unquantized (BF16/FP16/FP32) MoE
160 +# ===========================================================================
161 +
162 +
163 +class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
164 + """Portable vector grouped-gemm unquantized MoE experts."""
165 +
166 + isa = "vec"
167 + output_alignment = 32
168 + reduction_alignment = 1
169 +
170 + @classmethod
171 + def _intermediate_alignment(cls) -> int:
172 + return math.lcm(cls.output_alignment, cls.reduction_alignment)
173 +
174 + @classmethod
175 + def _padded_intermediate_size(cls, moe_config: FusedMoEConfig) -> int:
176 + intermediate_size = moe_config.intermediate_size_per_partition
177 + if moe_config.activation == MoEActivation.SWIGLUOAI:
178 + return intermediate_size
179 + return round_up(intermediate_size, cls._intermediate_alignment())
180 +
181 + @classmethod
182 + def _supports_grouped_gemm(
183 + cls,
184 + moe_config: FusedMoEConfig,
185 + ) -> tuple[bool, str | None]:
186 + intermediate_size = cls._padded_intermediate_size(moe_config)
187 + if (
188 + moe_config.hidden_dim % cls.output_alignment != 0
189 + or intermediate_size % cls.output_alignment != 0
190 + ):
191 + return False, (
192 + "kernel requires hidden and intermediate dimensions divisible by "
193 + f"{cls.output_alignment}"
194 + )
195 + if (
196 + moe_config.hidden_dim % cls.reduction_alignment != 0
197 + or intermediate_size % cls.reduction_alignment != 0
198 + ):
199 + return False, (
200 + "kernel requires reduction dimensions divisible by "
201 + f"{cls.reduction_alignment}"
202 + )
203 + return True, None
204 +
205 + def __init__(
206 + self,
207 + moe_config: FusedMoEConfig,
208 + quant_config: FusedMoEQuantConfig,
209 + ):
210 + super().__init__(moe_config, quant_config)
211 + # Router configuration that the monolithic apply() signature cannot
212 + # carry. Captured off the layer in process_weights_after_loading.
213 + self.use_grouped_topk = False
214 + self.renormalize = False
215 + self.scoring_func = "softmax"
216 + self.custom_routing_function: Callable | None = None
217 +
218 + @property
219 + def expects_unquantized_inputs(self) -> bool:
220 + return True
221 +
222 + @staticmethod
223 + def activation_format() -> mk.FusedMoEActivationFormat:
224 + return mk.FusedMoEActivationFormat.Standard
225 +
226 + @staticmethod
227 + def _supports_current_device() -> bool:
228 + return current_platform.is_cpu()
229 +
230 + @staticmethod
231 + def is_supported_config(
232 + cls: type[mk.FusedMoEExperts],
233 + moe_config: FusedMoEConfig,
234 + weight_key: QuantKey | None,
235 + activation_key: QuantKey | None,
236 + activation_format: mk.FusedMoEActivationFormat,
237 + ) -> tuple[bool, str | None]:
238 + supported, reason = mk.FusedMoEExperts.is_supported_config(
239 + cls, moe_config, weight_key, activation_key, activation_format
240 + )
241 + if not supported:
242 + return supported, reason
243 + cpu_cls = cast(type[CPUUnquantizedExperts], cls)
244 + return cpu_cls._supports_grouped_gemm(moe_config)
245 +
246 + @staticmethod
247 + def _supports_no_act_and_mul() -> bool:
248 + return False
249 +
250 + @staticmethod
251 + def _supports_activation(activation: MoEActivation) -> bool:
252 + return activation in (
253 + MoEActivation.SILU,
254 + MoEActivation.SWIGLUOAI,
255 + MoEActivation.GELU,
256 + MoEActivation.GELU_TANH,
257 + )
258 +
259 + @staticmethod
260 + def _supports_quant_scheme(
261 + weight_key: QuantKey | None,
262 + activation_key: QuantKey | None,
263 + ) -> bool:
264 + return (weight_key, activation_key) == (None, None)
265 +
266 + @staticmethod
267 + def _supports_parallel_config(
268 + moe_parallel_config: FusedMoEParallelConfig,
269 + ) -> bool:
270 + return True
271 +
272 + @staticmethod
273 + def _supports_routing_method(
274 + routing_method: RoutingMethodType,
275 + weight_key: QuantKey | None,
276 + activation_key: QuantKey | None,
277 + ) -> bool:
278 + # Routing runs in select_experts(), which covers every routing method
279 + # a layer can be configured with, including custom routing functions.
280 + return True
281 +
282 + @staticmethod
283 + def _supports_router_logits_dtype(
284 + router_logits_dtype: torch.dtype | None,
285 + routing_method: RoutingMethodType,
286 + ) -> bool:
287 + return True
288 +
289 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
290 + self._pad_moe_intermediate(layer)
291 + self.use_grouped_topk = layer.use_grouped_topk
292 + self.renormalize = layer.renormalize
293 + self.scoring_func = layer.scoring_func
294 + self.custom_routing_function = layer.custom_routing_function
295 + replace_parameter(
296 + layer, "w13_weight", cpu_prepack_moe_weight(layer.w13_weight, self.isa)
297 + )
298 + replace_parameter(
299 + layer, "w2_weight", cpu_prepack_moe_weight(layer.w2_weight, self.isa)
300 + )
301 +
302 + def _pad_moe_intermediate(self, layer: torch.nn.Module) -> None:
303 + """Zero-pad the per-partition MoE intermediate dim of both weights and
304 + the expert bias, see `_padded_intermediate_size`."""
305 + intermediate_size = self.moe_config.intermediate_size_per_partition
306 + padded_size = self._padded_intermediate_size(self.moe_config)
307 + if padded_size == intermediate_size:
308 + return
309 +
310 + num_experts, _, hidden_size = layer.w13_weight.shape
311 +
312 + new_w13 = layer.w13_weight.new_zeros(num_experts, 2 * padded_size, hidden_size)
313 + new_w13[:, :intermediate_size] = layer.w13_weight[:, :intermediate_size]
314 + new_w13[:, padded_size : padded_size + intermediate_size] = layer.w13_weight[
315 + :, intermediate_size:
316 + ]
317 + replace_parameter(layer, "w13_weight", new_w13)
318 +
319 + new_w2 = layer.w2_weight.new_zeros(num_experts, hidden_size, padded_size)
320 + new_w2[:, :, :intermediate_size] = layer.w2_weight
321 + replace_parameter(layer, "w2_weight", new_w2)
322 +
323 + if hasattr(layer, "w13_bias"):
324 + new_bias = layer.w13_bias.new_zeros(num_experts, 2 * padded_size)
325 + new_bias[:, :intermediate_size] = layer.w13_bias[:, :intermediate_size]
326 + new_bias[:, padded_size : padded_size + intermediate_size] = layer.w13_bias[
327 + :, intermediate_size:
328 + ]
329 + # Assign through .data rather than replacing the Parameter: the
330 + # quant config is built before this runs and holds a reference to
331 + # this very object, which is what feeds self.w1_bias in apply().
332 + layer.w13_bias.data = new_bias
333 +
334 + def _select_experts(
335 + self,
336 + hidden_states: torch.Tensor,
337 + router_logits: torch.Tensor,
338 + num_expert_group: int | None,
339 + topk_group: int | None,
340 + e_score_correction_bias: torch.Tensor | None,
341 + routed_scaling_factor: float | None,
342 + ) -> tuple[torch.Tensor, torch.Tensor]:
343 + return select_experts(
344 + hidden_states=hidden_states,
345 + router_logits=router_logits,
346 + top_k=self.moe_config.experts_per_token,
347 + use_grouped_topk=self.use_grouped_topk,
348 + renormalize=self.renormalize,
349 + topk_group=topk_group,
350 + num_expert_group=num_expert_group,
351 + custom_routing_function=self.custom_routing_function,
352 + scoring_func=self.scoring_func,
353 + routed_scaling_factor=(
354 + routed_scaling_factor if routed_scaling_factor is not None else 1.0
355 + ),
356 + e_score_correction_bias=e_score_correction_bias,
357 + )
358 +
359 + def apply(
360 + self,
361 + hidden_states: torch.Tensor,
362 + w1: torch.Tensor,
363 + w2: torch.Tensor,
364 + router_logits: torch.Tensor,
365 + activation: MoEActivation,
366 + global_num_experts: int,
367 + expert_map: torch.Tensor | None,
368 + a1q_scale: torch.Tensor | None,
369 + apply_router_weight_on_input: bool,
370 + # grouped topk + fused topk bias parameters
371 + num_expert_group: int | None = None,
372 + e_score_correction_bias: torch.Tensor | None = None,
373 + routed_scaling_factor: float | None = None,
374 + topk_group: int | None = None,
375 + ) -> torch.Tensor:
376 + topk_weights, topk_ids = self._select_experts(
377 + hidden_states,
378 + router_logits,
379 + num_expert_group,
380 + topk_group,
381 + e_score_correction_bias,
382 + routed_scaling_factor,
383 + )
384 +
385 + if apply_router_weight_on_input:
386 + assert topk_ids.size(1) == 1, (
387 + "apply_router_weight_on_input is only implemented for topk=1"
388 + )
389 + hidden_states.mul_(topk_weights.to(hidden_states.dtype))
390 +
391 + return cpu_fused_moe(
392 + hidden_states,
393 + w1,
394 + w2,
395 + self.w1_bias,
396 + self.w2_bias,
397 + topk_weights,
398 + topk_ids,
399 + activation.value,
400 + self.isa,
401 + apply_router_weight_on_input,
402 + )
403 +
404 +
405 +class X86CPUUnquantizedExperts(CPUUnquantizedExperts):
406 + """x86 AMX grouped-gemm unquantized MoE experts."""
407 +
408 + isa = "amx"
409 + output_alignment = 32
410 + reduction_alignment = 32
411 +
412 + @staticmethod
413 + def _supports_current_device() -> bool:
414 + return (
415 + current_platform.is_cpu()
416 + and current_platform.get_cpu_architecture() == CpuArchEnum.X86
417 + and torch.cpu._is_amx_tile_supported()
418 + )
419 +
420 + @staticmethod
421 + def is_supported_config(
422 + cls: type[mk.FusedMoEExperts],
423 + moe_config: FusedMoEConfig,
424 + weight_key: QuantKey | None,
425 + activation_key: QuantKey | None,
426 + activation_format: mk.FusedMoEActivationFormat,
427 + ) -> tuple[bool, str | None]:
428 + supported, reason = mk.FusedMoEExperts.is_supported_config(
429 + cls, moe_config, weight_key, activation_key, activation_format
430 + )
431 + if not supported:
432 + return supported, reason
433 + if moe_config.in_dtype != torch.bfloat16:
434 + return False, "kernel requires bfloat16 activations"
435 + cpu_cls = cast(type[CPUUnquantizedExperts], cls)
436 + return cpu_cls._supports_grouped_gemm(moe_config)
437 +
438 +
439 +class ArmCPUUnquantizedExperts(CPUUnquantizedExperts):
440 + """Arm NEON grouped-gemm unquantized MoE experts."""
441 +
442 + isa = "neon"
443 + output_alignment = 32
444 + reduction_alignment = 4
445 +
446 + @staticmethod
447 + def _supports_current_device() -> bool:
448 + return (
449 + current_platform.is_cpu()
450 + and current_platform.get_cpu_architecture() == CpuArchEnum.ARM
451 + and sys.platform != "darwin"
452 + )
453 +
454 + @staticmethod
455 + def is_supported_config(
456 + cls: type[mk.FusedMoEExperts],
457 + moe_config: FusedMoEConfig,
458 + weight_key: QuantKey | None,
459 + activation_key: QuantKey | None,
460 + activation_format: mk.FusedMoEActivationFormat,
461 + ) -> tuple[bool, str | None]:
462 + supported, reason = mk.FusedMoEExperts.is_supported_config(
463 + cls, moe_config, weight_key, activation_key, activation_format
464 + )
465 + if not supported:
466 + return supported, reason
467 + if moe_config.in_dtype != torch.bfloat16:
468 + return False, "kernel requires bfloat16 activations"
469 + cpu_cls = cast(type[CPUUnquantizedExperts], cls)
470 + return cpu_cls._supports_grouped_gemm(moe_config)
471 +
472
473 # ===========================================================================
474 # FP8 W8A16 MoE
@@ -131,10 +570,6 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic):
570 routed_scaling_factor: float | None = None,
571 topk_group: int | None = None,
572 ) -> torch.Tensor:
134 - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
135 - select_experts,
136 - )
137 -
573 topk_weights, topk_ids = select_experts(
574 hidden_states=hidden_states,
575 router_logits=router_logits,
@@ -289,10 +724,6 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic):
724 routed_scaling_factor: float | None = None,
725 topk_group: int | None = None,
726 ) -> torch.Tensor:
292 - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
293 - select_experts,
294 - )
295 -
727 topk_weights, topk_ids = select_experts(
728 hidden_states=hidden_states,
729 router_logits=router_logits,
@@ -505,10 +936,6 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
936 "apply_router_weight_on_input=True. "
937 )
938
508 - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
509 - select_experts,
510 - )
511 -
939 topk_weights, topk_ids = select_experts(
940 hidden_states=hidden_states,
941 router_logits=router_logits,
@@ -631,7 +1058,6 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1058
1059 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1060 """VNNI-prepack INT8 MoE weights for CPU kernel."""
634 - from vllm.model_executor.utils import replace_parameter
1061
1062 w13 = torch.ops._C.convert_weight_packed(layer.w13_weight)
1063 w2 = torch.ops._C.convert_weight_packed(layer.w2_weight)
@@ -655,10 +1081,6 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1081 routed_scaling_factor: float | None = None,
1082 topk_group: int | None = None,
1083 ) -> torch.Tensor:
658 - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
659 - select_experts,
660 - )
661 -
1084 topk_weights, topk_ids = select_experts(
1085 hidden_states=hidden_states,
1086 router_logits=router_logits,
@@ -796,8 +1218,6 @@ class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1218 return True
1219
1220 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
799 - from vllm.model_executor.utils import replace_parameter
800 -
1221 w13 = cpu_prepack_moe_weight_int8(layer.w13_weight, "neon")
1222 w2 = cpu_prepack_moe_weight_int8(layer.w2_weight, "neon")
1223 replace_parameter(layer, "w13_weight", w13)
@@ -819,10 +1239,6 @@ class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1239 routed_scaling_factor: float | None = None,
1240 topk_group: int | None = None,
1241 ) -> torch.Tensor:
822 - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
823 - select_experts,
824 - )
825 -
1242 topk_weights, topk_ids = select_experts(
1243 hidden_states=hidden_states,
1244 router_logits=router_logits,
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+14 -4
@@ -143,6 +143,20 @@ def backend_to_kernel_cls(
143
144 return [XPUExperts]
145
146 + elif backend == UnquantizedMoeBackend.CPU:
147 + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
148 + ArmCPUUnquantizedExperts,
149 + CPUUnquantizedExperts,
150 + X86CPUUnquantizedExperts,
151 + )
152 +
153 + # Prefer architecture-specific kernels before the portable vector path.
154 + return [
155 + X86CPUUnquantizedExperts,
156 + ArmCPUUnquantizedExperts,
157 + CPUUnquantizedExperts,
158 + ]
159 +
160 else:
161 raise ValueError(f"Unknown unquantized MoE backend: {backend.value}")
162
@@ -199,10 +213,6 @@ def select_unquantized_moe_backend(
213 Note: Shape-specific fallbacks may still occur at runtime.
214 """
215
202 - if current_platform.is_cpu():
203 - # TODO: migrate to MK structure.
204 - return UnquantizedMoeBackend.CPU, None
205 -
216 if current_platform.is_tpu():
217 return UnquantizedMoeBackend.TPU, None
218
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+24 -79
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from collections.abc import Callable
4 from typing import TYPE_CHECKING
5
6 import torch
@@ -33,7 +32,6 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
32 )
33 from vllm.model_executor.utils import replace_parameter, set_weight_attrs
34 from vllm.platforms import current_platform
36 -from vllm.platforms.interface import CpuArchEnum
35
36 if TYPE_CHECKING:
37 from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
@@ -54,13 +52,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
52 moe_config=self.moe,
53 )
54
57 - @property
58 - def is_monolithic(self) -> bool:
59 - # Escape hatch for CPU, which stays on the old monolithic path.
60 - if self.unquantized_backend == UnquantizedMoeBackend.CPU:
61 - return True
62 - return super().is_monolithic
63 -
55 @property
56 def supports_eplb(self) -> bool:
57 return True
@@ -71,8 +62,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
62 ):
63 raise ValueError(
64 f"{self.__class__.__name__} uses the new modular kernel initialization "
74 - "logic for all but the CPU backend. CPU backend is monolithic. "
75 - "So this function should not be called."
65 + "logic. So this function should not be called."
66 )
67
68 def select_gemm_impl(
@@ -201,6 +191,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
191 routing_tables=layer._expert_routing_tables(),
192 )
193
194 + if self.unquantized_backend == UnquantizedMoeBackend.CPU:
195 + # The CPU experts need the layer itself for the setup that
196 + # convert_to_unquantized_kernel_format cannot express, since
197 + # it only sees the two weight tensors: padding and prepacking
198 + # into the grouped-gemm layout (bias included), and capturing
199 + # the router config that monolithic apply() cannot carry.
200 + self.moe_kernel.fused_experts.process_weights_after_loading(layer)
201 +
202 def process_weights_after_loading(self, layer: "RoutedExperts") -> None:
203 super().process_weights_after_loading(layer)
204
@@ -221,38 +219,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
219 # OOT handles internally.
220 return
221
224 - elif self.unquantized_backend == UnquantizedMoeBackend.CPU:
225 - # CPU stays on the old path — no oracle, no moe_kernel.
226 - from vllm.model_executor.layers.fused_moe import cpu_fused_moe
227 -
228 - if current_platform.get_cpu_architecture() == CpuArchEnum.X86:
229 - from vllm.model_executor.layers.utils import check_cpu_sgl_kernel
230 -
231 - dtype_w13 = layer.w13_weight.dtype
232 - _, n_w13, k_w13 = layer.w13_weight.size()
233 - dtype_w2 = layer.w2_weight.dtype
234 - _, n_w2, k_w2 = layer.w2_weight.size()
235 - if (
236 - envs.VLLM_CPU_SGL_KERNEL
237 - and check_cpu_sgl_kernel(n_w13, k_w13, dtype_w13)
238 - and check_cpu_sgl_kernel(n_w2, k_w2, dtype_w2)
239 - ):
240 - packed_w13_weight = torch.ops._C.convert_weight_packed(
241 - layer.w13_weight
242 - )
243 - assert packed_w13_weight.size() == layer.w13_weight.size()
244 - layer.w13_weight.copy_(packed_w13_weight)
245 - del packed_w13_weight
246 - packed_w2_weight = torch.ops._C.convert_weight_packed(
247 - layer.w2_weight
248 - )
249 - assert packed_w2_weight.size() == layer.w2_weight.size()
250 - layer.w2_weight.copy_(packed_w2_weight)
251 - self.cpu_fused_moe: Callable = cpu_fused_moe.SGLFusedMOE(layer)
252 - else:
253 - self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
254 - else:
255 - self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
222 elif self.unquantized_backend == UnquantizedMoeBackend.XPU:
223 w13 = layer.w13_weight
224 w2 = layer.w2_weight
@@ -363,39 +329,18 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
329 input_ids: torch.Tensor | None = None,
330 ) -> torch.Tensor:
331 assert self.is_monolithic
366 - if self.unquantized_backend == UnquantizedMoeBackend.CPU:
367 - assert self.moe_kernel is None
368 - return self.cpu_fused_moe(
369 - layer,
370 - x,
371 - layer.use_grouped_topk,
372 - layer.top_k,
373 - router_logits,
374 - layer.renormalize,
375 - layer.topk_group,
376 - layer.num_expert_group,
377 - layer.global_num_experts,
378 - layer.expert_map,
379 - layer.custom_routing_function,
380 - layer.scoring_func,
381 - layer.routed_scaling_factor,
382 - layer.e_score_correction_bias,
383 - layer.apply_router_weight_on_input,
384 - layer.activation,
385 - )
386 - else:
387 - assert self.moe_kernel is not None
388 - return self.moe_kernel.apply_monolithic(
389 - x,
390 - layer.w13_weight,
391 - layer.w2_weight,
392 - router_logits,
393 - activation=layer.activation,
394 - global_num_experts=layer.global_num_experts,
395 - expert_map=layer.expert_map,
396 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
397 - num_expert_group=layer.num_expert_group,
398 - topk_group=layer.topk_group,
399 - e_score_correction_bias=layer.e_score_correction_bias,
400 - routed_scaling_factor=layer.routed_scaling_factor,
401 - )
332 + assert self.moe_kernel is not None
333 + return self.moe_kernel.apply_monolithic(
334 + x,
335 + layer.w13_weight,
336 + layer.w2_weight,
337 + router_logits,
338 + activation=layer.activation,
339 + global_num_experts=layer.global_num_experts,
340 + expert_map=layer.expert_map,
341 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
342 + num_expert_group=layer.num_expert_group,
343 + topk_group=layer.topk_group,
344 + e_score_correction_bias=layer.e_score_correction_bias,
345 + routed_scaling_factor=layer.routed_scaling_factor,
346 + )