From 0a6446005d51c9e6bfa09352f7f288ddeff17c77 Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Mon, 3 Aug 2026 14:55:04 +0800 Subject: [PATCH] [CPU] Migrate unquantized MoE to the modular-kernel experts structure (#50133) Signed-off-by: jiang1.li Signed-off-by: Fadi Arafeh Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Fadi Arafeh --- cmake/cpu_extension.cmake | 3 +- csrc/cpu/cpu_fused_moe_activations.hpp | 37 +- csrc/cpu/cpu_types_riscv_impl.hpp | 22 + csrc/cpu/cpu_types_scalar.hpp | 44 ++ csrc/cpu/cpu_types_vsx.hpp | 23 + csrc/cpu/cpu_types_vxe.hpp | 66 +- csrc/cpu/cpu_types_x86.hpp | 35 ++ csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp | 3 +- csrc/cpu/torch_bindings.cpp | 3 - csrc/cpu/utils.hpp | 2 - docs/design/moe_kernel_features.md | 2 +- docs/getting_started/installation/cpu.md | 2 +- tests/kernels/moe/test_cpu_fused_moe.py | 209 +++++-- tests/kernels/moe/test_moe.py | 79 --- .../moe/test_unquantized_backend_selection.py | 117 +++- .../layers/fused_moe/cpu_fused_moe.py | 577 ------------------ .../layers/fused_moe/experts/cpu_int4_moe.py | 2 +- .../layers/fused_moe/experts/cpu_moe.py | 464 +++++++++++++- .../layers/fused_moe/oracle/unquantized.py | 18 +- .../fused_moe/unquantized_fused_moe_method.py | 103 +--- 20 files changed, 972 insertions(+), 839 deletions(-) delete mode 100644 vllm/model_executor/layers/fused_moe/cpu_fused_moe.py diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index cdfc6cb6f84..c66854386c4 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -461,6 +461,7 @@ set(VLLM_EXT_SRC "csrc/cpu/pos_encoding.cpp" "csrc/cpu/mamba_cpu.cpp" "csrc/moe/dynamic_4bit_int_moe_cpu.cpp" + "csrc/cpu/cpu_fused_moe.cpp" "csrc/cpu/cpu_attn.cpp" "csrc/cpu/torch_bindings.cpp") @@ -478,7 +479,6 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) "csrc/cpu/cpu_tanhf_neon.hpp" ${VLLM_EXT_SRC}) if (ARM_BF16_FOUND) - set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe.cpp" ${VLLM_EXT_SRC}) if (ARM_I8MM_FOUND) set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe_int8.cpp" ${VLLM_EXT_SRC}) endif() @@ -535,6 +535,7 @@ if (ENABLE_X86_ISA) set(VLLM_EXT_SRC_AVX2 "csrc/cpu/sgl-kernels/fla.cpp" + "csrc/cpu/cpu_fused_moe.cpp" "csrc/cpu/utils.cpp" "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/cpu_attn.cpp" diff --git a/csrc/cpu/cpu_fused_moe_activations.hpp b/csrc/cpu/cpu_fused_moe_activations.hpp index 31e8cb009b7..a5570cee97d 100644 --- a/csrc/cpu/cpu_fused_moe_activations.hpp +++ b/csrc/cpu/cpu_fused_moe_activations.hpp @@ -11,6 +11,13 @@ #include "cpu/cpu_arch_macros.h" #include "cpu/utils.hpp" +#if defined(DEFINE_FAST_EXP) + #define DEFINE_CPU_FUSED_MOE_EXP DEFINE_FAST_EXP +#else + #define DEFINE_CPU_FUSED_MOE_EXP \ + auto fast_exp = [](const vec_op::FP32Vec16& vec) { return vec.exp(); }; +#endif + namespace cpu_fused_moe_utils { enum class FusedMOEAct { SiluAndMul, @@ -39,7 +46,7 @@ void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, const int32_t input_stride, const int32_t output_stride) { using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t; -#if !defined(__aarch64__) +#if defined(__AVX512F__) // For GPT-OSS interleaved gate-up weights alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30}; @@ -50,7 +57,7 @@ void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, vec_op::FP32Vec16 alpha_vec(1.702); vec_op::FP32Vec16 one_vec(1.0); - DEFINE_FAST_EXP + DEFINE_CPU_FUSED_MOE_EXP for (int32_t m = 0; m < m_size; ++m) { 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, vec_op::FP32Vec16 gate_vec(vec_op::uninit); vec_op::FP32Vec16 up_vec(vec_op::uninit); vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec); -#else +#elif defined(__AVX512F__) vec_op::FP32Vec16 gate_vec(input + n, index_vec); vec_op::FP32Vec16 up_vec(input + n + 1, index_vec); +#else + alignas(64) float gate_values[16]; + alignas(64) float up_values[16]; + const float* interleaved = input + n; + for (int32_t i = 0; i < 16; ++i) { + gate_values[i] = interleaved[2 * i]; + up_values[i] = interleaved[2 * i + 1]; + } + vec_op::FP32Vec16 gate_vec(gate_values); + vec_op::FP32Vec16 up_vec(up_values); #endif gate_vec = gate_vec.min(gate_up_max_vec); 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, float* __restrict__ up = input + dim; vec_op::FP32Vec16 one_vec(1.0); - DEFINE_FAST_EXP + DEFINE_CPU_FUSED_MOE_EXP for (int32_t m = 0; m < m_size; ++m) { for (int32_t n = 0; n < dim; n += 16) { @@ -115,21 +132,13 @@ void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, vec_op::FP32Vec16 one_vec(1.0); vec_op::FP32Vec16 w1_vec(M_SQRT1_2); vec_op::FP32Vec16 w2_vec(0.5); - alignas(64) float temp[16]; - - DEFINE_FAST_EXP for (int32_t m = 0; m < m_size; ++m) { for (int32_t n = 0; n < dim; n += 16) { vec_op::FP32Vec16 gate_vec(gate + n); vec_op::FP32Vec16 up_vec(up + n); auto er_input_vec = gate_vec * w1_vec; - - er_input_vec.save(temp); - for (int32_t i = 0; i < 16; ++i) { - temp[i] = std::erf(temp[i]); - } - vec_op::FP32Vec16 er_vec(temp); + auto er_vec = er_input_vec.er(); auto gelu = gate_vec * w2_vec * (one_vec + er_vec); auto gated_output_fp32 = up_vec * gelu; scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); @@ -201,4 +210,6 @@ FORCE_INLINE void apply_gated_act(const FusedMOEAct act, } } // namespace cpu_fused_moe_utils +#undef DEFINE_CPU_FUSED_MOE_EXP + #endif diff --git a/csrc/cpu/cpu_types_riscv_impl.hpp b/csrc/cpu/cpu_types_riscv_impl.hpp index 70cb0ab52de..106bf44a241 100644 --- a/csrc/cpu/cpu_types_riscv_impl.hpp +++ b/csrc/cpu/cpu_types_riscv_impl.hpp @@ -678,6 +678,9 @@ struct FP32Vec16 : public Vec { return FP32Vec16( RVVI(__riscv_vfsub_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); } + FP32Vec16 operator-() const { + return FP32Vec16(RVVI(__riscv_vfneg_v_f32, LMUL_512)(reg, VEC_ELEM_NUM)); + } FP32Vec16 operator*(const FP32Vec16& b) const { return FP32Vec16( RVVI(__riscv_vfmul_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); @@ -898,6 +901,25 @@ struct INT8Vec16 : public Vec { } }; +// Reference implementation for vector operations missing from some backends. +struct INT8Vec64 { + constexpr static int VEC_ELEM_NUM = 64; + + explicit INT8Vec64(const int8_t* ptr) { + std::memcpy(data_, ptr, sizeof(data_)); + } + + void save(int8_t* ptr) const { std::memcpy(ptr, data_, sizeof(data_)); } + + void save(int8_t* ptr, const int elem_num) const { + TORCH_CHECK(elem_num > 0 && elem_num <= VEC_ELEM_NUM); + std::memcpy(ptr, data_, elem_num); + } + + private: + int8_t data_[VEC_ELEM_NUM]; +}; + // ============================================================================ // Type Traits & Global Helpers // ============================================================================ diff --git a/csrc/cpu/cpu_types_scalar.hpp b/csrc/cpu/cpu_types_scalar.hpp index 94b5179b171..d94cdbc076e 100644 --- a/csrc/cpu/cpu_types_scalar.hpp +++ b/csrc/cpu/cpu_types_scalar.hpp @@ -333,6 +333,13 @@ struct FP32Vec16 : public Vec { return FP32Vec16(ret); } + FP32Vec16 operator-() const { + f32x16_t ret; + unroll_loop( + [&ret, this](int i) { ret.val[i] = -reg.val[i]; }); + return FP32Vec16(ret); + } + FP32Vec16 operator/(const FP32Vec16& b) const { f32x16_t ret; unroll_loop( @@ -356,6 +363,10 @@ struct FP32Vec16 : public Vec { return FP32Vec16(ret); } + FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const { + return this->max(min_v).min(max_v); + } + FP32Vec16 abs() const { f32x16_t ret; unroll_loop( @@ -363,6 +374,20 @@ struct FP32Vec16 : public Vec { return FP32Vec16(ret); } + FP32Vec16 exp() const { + f32x16_t ret; + unroll_loop( + [&ret, this](int i) { ret.val[i] = std::exp(reg.val[i]); }); + return FP32Vec16(ret); + } + + FP32Vec16 er() const { + f32x16_t ret; + unroll_loop( + [&ret, this](int i) { ret.val[i] = std::erf(reg.val[i]); }); + return FP32Vec16(ret); + } + FP32Vec16 tanh() const { f32x16_t ret; unroll_loop( @@ -404,6 +429,25 @@ struct FP32Vec16 : public Vec { void save(void* ptr) const { *reinterpret_cast(ptr) = reg; } }; +// Reference implementation for vector operations missing from some backends. +struct INT8Vec64 { + constexpr static int VEC_ELEM_NUM = 64; + + explicit INT8Vec64(const int8_t* ptr) { + std::memcpy(data_, ptr, sizeof(data_)); + } + + void save(int8_t* ptr) const { std::memcpy(ptr, data_, sizeof(data_)); } + + void save(int8_t* ptr, const int elem_num) const { + TORCH_CHECK(elem_num > 0 && elem_num <= VEC_ELEM_NUM); + std::memcpy(ptr, data_, elem_num); + } + + private: + int8_t data_[VEC_ELEM_NUM]; +}; + template struct VecType { using vec_type = void; diff --git a/csrc/cpu/cpu_types_vsx.hpp b/csrc/cpu/cpu_types_vsx.hpp index 42083bc3eb6..d46906130e9 100644 --- a/csrc/cpu/cpu_types_vsx.hpp +++ b/csrc/cpu/cpu_types_vsx.hpp @@ -655,6 +655,11 @@ struct FP32Vec16 : public Vec { vec_sub(reg.val[3], b.reg.val[3])})); } + FP32Vec16 operator-() const { + return FP32Vec16( + f32x4x4_t({-reg.val[0], -reg.val[1], -reg.val[2], -reg.val[3]})); + } + FP32Vec16 operator/(const FP32Vec16& b) const { return FP32Vec16(f32x4x4_t({vec_div(reg.val[0], b.reg.val[0]), vec_div(reg.val[1], b.reg.val[1]), @@ -756,6 +761,24 @@ struct FP32Vec16 : public Vec { hi_e.reg.val[0], hi_e.reg.val[1]}); } + FP32Vec16 tanh() const { + FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]}); + FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]}); + auto lo_tanh = lo.tanh(); + auto hi_tanh = hi.tanh(); + return FP32Vec16(f32x4x4_t{lo_tanh.reg.val[0], lo_tanh.reg.val[1], + hi_tanh.reg.val[0], hi_tanh.reg.val[1]}); + } + + FP32Vec16 er() const { + FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]}); + FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]}); + auto lo_er = lo.er(); + auto hi_er = hi.er(); + return FP32Vec16(f32x4x4_t{lo_er.reg.val[0], lo_er.reg.val[1], + hi_er.reg.val[0], hi_er.reg.val[1]}); + } + float reduce_max() { __vector float max01 = vec_max(reg.val[0], reg.val[1]); __vector float max23 = vec_max(reg.val[2], reg.val[3]); diff --git a/csrc/cpu/cpu_types_vxe.hpp b/csrc/cpu/cpu_types_vxe.hpp index 9b28300cedf..f705213a60a 100644 --- a/csrc/cpu/cpu_types_vxe.hpp +++ b/csrc/cpu/cpu_types_vxe.hpp @@ -6,8 +6,10 @@ #include #include #include +#include #include #include + namespace vec_op { struct fp8_e4m3_tag {}; @@ -727,6 +729,11 @@ struct FP32Vec16 : public Vec { vec_sub(reg.val[3], b.reg.val[3])})); } + FP32Vec16 operator-() const { + return FP32Vec16(f32x4x4_t({vec_neg(reg.val[0]), vec_neg(reg.val[1]), + vec_neg(reg.val[2]), vec_neg(reg.val[3])})); + } + FP32Vec16 operator/(const FP32Vec16& b) const { return FP32Vec16(f32x4x4_t({vec_div(reg.val[0], b.reg.val[0]), vec_div(reg.val[1], b.reg.val[1]), @@ -734,6 +741,33 @@ struct FP32Vec16 : public Vec { vec_div(reg.val[3], b.reg.val[3])})); } + FP32Vec16 exp() const { + FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]}); + FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]}); + auto lo_exp = lo.exp(); + auto hi_exp = hi.exp(); + return FP32Vec16(f32x4x4_t{lo_exp.reg.val[0], lo_exp.reg.val[1], + hi_exp.reg.val[0], hi_exp.reg.val[1]}); + } + + FP32Vec16 tanh() const { + FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]}); + FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]}); + auto lo_tanh = lo.tanh(); + auto hi_tanh = hi.tanh(); + return FP32Vec16(f32x4x4_t{lo_tanh.reg.val[0], lo_tanh.reg.val[1], + hi_tanh.reg.val[0], hi_tanh.reg.val[1]}); + } + + FP32Vec16 er() const { + FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]}); + FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]}); + auto lo_er = lo.er(); + auto hi_er = hi.er(); + return FP32Vec16(f32x4x4_t{lo_er.reg.val[0], lo_er.reg.val[1], + hi_er.reg.val[0], hi_er.reg.val[1]}); + } + float reduce_sum() const { __vector float sum = vec_add(vec_add(reg.val[0], reg.val[1]), vec_add(reg.val[2], reg.val[3])); @@ -765,6 +799,17 @@ struct FP32Vec16 : public Vec { vec_max(reg.val[3], b.reg.val[3])})); } + FP32Vec16 min(const FP32Vec16& b) const { + return FP32Vec16(f32x4x4_t({vec_min(reg.val[0], b.reg.val[0]), + vec_min(reg.val[1], b.reg.val[1]), + vec_min(reg.val[2], b.reg.val[2]), + vec_min(reg.val[3], b.reg.val[3])})); + } + + FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const { + return this->max(min_v).min(max_v); + } + float reduce_max() const { __vector float m = vec_max(vec_max(reg.val[0], reg.val[1]), vec_max(reg.val[2], reg.val[3])); @@ -899,6 +944,25 @@ struct INT8Vec16 : public Vec { } }; +// Reference implementation for vector operations missing from some backends. +struct INT8Vec64 { + constexpr static int VEC_ELEM_NUM = 64; + + explicit INT8Vec64(const int8_t* ptr) { + std::memcpy(data_, ptr, sizeof(data_)); + } + + void save(int8_t* ptr) const { std::memcpy(ptr, data_, sizeof(data_)); } + + void save(int8_t* ptr, const int elem_num) const { + TORCH_CHECK(elem_num > 0 && elem_num <= VEC_ELEM_NUM); + std::memcpy(ptr, data_, elem_num); + } + + private: + int8_t data_[VEC_ELEM_NUM]; +}; + template struct VecType { using vec_type = void; @@ -1325,4 +1389,4 @@ FORCE_INLINE void prefetch(const void* addr) { }; // namespace vec_op -#endif \ No newline at end of file +#endif diff --git a/csrc/cpu/cpu_types_x86.hpp b/csrc/cpu/cpu_types_x86.hpp index d2a72ce9ccd..720221d72c8 100644 --- a/csrc/cpu/cpu_types_x86.hpp +++ b/csrc/cpu/cpu_types_x86.hpp @@ -2,6 +2,7 @@ #ifndef CPU_TYPES_X86_HPP #define CPU_TYPES_X86_HPP +#include #include #include #include @@ -595,6 +596,8 @@ struct FP32Vec16 : public Vec { FP32Vec16 tanh() const { return FP32Vec16(Sleef_tanhf16_u10(reg)); } + FP32Vec16 er() const { return FP32Vec16(Sleef_erff16_u10(reg)); } + float reduce_sum() const { return _mm512_reduce_add_ps(reg); } float reduce_max() const { return _mm512_reduce_max_ps(reg); } @@ -798,6 +801,14 @@ struct FP32Vec16 : public Vec { return FP32Vec16(low.tanh().reg, high.tanh().reg); } + FP32Vec16 exp() const { + return FP32Vec16(Sleef_expf8_u10(reg_low), Sleef_expf8_u10(reg_high)); + } + + FP32Vec16 er() const { + return FP32Vec16(Sleef_erff8_u10(reg_low), Sleef_erff8_u10(reg_high)); + } + FP32Vec16 min(const FP32Vec16& b) const { return FP32Vec16(_mm256_min_ps(reg_low, b.reg_low), _mm256_min_ps(reg_high, b.reg_high)); @@ -927,6 +938,30 @@ struct INT8Vec16 : public Vec { for (int i = 0; i < elem_num; ++i) ptr[i] = ar.values[i]; } }; + +struct INT8Vec64 : public Vec { + constexpr static int VEC_ELEM_NUM = 64; + + __m256i reg_low; + __m256i reg_high; + + explicit INT8Vec64(const void* ptr) + : reg_low(_mm256_loadu_si256(reinterpret_cast(ptr))), + reg_high( + _mm256_loadu_si256(reinterpret_cast(ptr) + 1)) {} + + void save(void* ptr) const { + _mm256_storeu_si256(reinterpret_cast<__m256i*>(ptr), reg_low); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(ptr) + 1, reg_high); + } + + void save(int8_t* ptr, const int elem_num) const { + TORCH_CHECK(elem_num > 0 && elem_num <= VEC_ELEM_NUM); + int8_t values[VEC_ELEM_NUM]; + save(values); + std::memcpy(ptr, values, elem_num); + } +}; #endif template diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp index ad7d4be113e..85330094aef 100644 --- a/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp @@ -71,8 +71,7 @@ class TileGemm82 { scalar_t* __restrict__ curr_m_a = curr_a; vec_op::unroll_loop([&](int32_t i) { scalar_t v = *curr_m_a; - load_vec_t a_reg_original(v); - vec_op::FP32Vec16 a_reg(a_reg_original); + vec_op::FP32Vec16 a_reg(static_cast(v)); c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg; c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg; diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index e6130473e30..7527ac4846a 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -614,7 +614,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { #endif // fused moe -#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) && !defined(__APPLE__)) ops.def( "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) " "-> ()"); @@ -625,8 +624,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "bool skip_weighted, " "str act, str isa) -> ()"); ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe); -#endif // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) && - // !defined(__APPLE__)) #if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \ !defined(__APPLE__) ops.def( diff --git a/csrc/cpu/utils.hpp b/csrc/cpu/utils.hpp index 78ee7081b24..23216f69e33 100644 --- a/csrc/cpu/utils.hpp +++ b/csrc/cpu/utils.hpp @@ -40,12 +40,10 @@ struct VecTypeTrait { using vec_t = vec_op::BF16Vec16; }; -#if !defined(__powerpc__) template <> struct VecTypeTrait { using vec_t = vec_op::FP16Vec16; }; -#endif struct Counter { std::atomic counter; diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 42b05e3c306..53f91de55e0 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -91,7 +91,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k | trtllm | standard | mxfp4,
nvfp4 | G(16),G(32) | 5 | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],
[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],
[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],
[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] | | hpc | standard | fp8 | G(128),T | silu | Y | Y | [`HPCExperts`][vllm.model_executor.layers.fused_moe.hpc_moe.HPCExperts] | | rocm aiter moe | standard | mxfp4,
fp8 | G(32),G(128),A,T | silu, gelu,
swigluoai | Y | N | `rocm_aiter_fused_experts`,
`AiterExperts` | -| cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] | +| cpu_moe | standard | N/A | N/A | silu, gelu,
gelu_tanh,
swigluoai | Y | N | [`X86CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.X86CPUUnquantizedExperts],
[`ArmCPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.ArmCPUUnquantizedExperts],
[`CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.CPUUnquantizedExperts] | | naive batched4 | batched | int8,
fp8 | G,A,T | silu, gelu | 6 | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] | !!! info "Table key" diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md index 925f54dcec7..3e6227b8d40 100644 --- a/docs/getting_started/installation/cpu.md +++ b/docs/getting_started/installation/cpu.md @@ -152,7 +152,7 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cpu VLLM_TARGET_DEVICE=cpu - `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. - `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`. - `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. -- `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). +- `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. - `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). ## FAQ diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py index 9e99a274fb3..e9dc1c48200 100644 --- a/tests/kernels/moe/test_cpu_fused_moe.py +++ b/tests/kernels/moe/test_cpu_fused_moe.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import sys + import pytest import torch +import vllm.model_executor.layers.fused_moe.modular_kernel as mk from tests.kernels.allclose_default import get_default_atol, get_default_rtol from vllm._custom_ops import ( cpu_fused_moe, @@ -12,9 +15,18 @@ from vllm._custom_ops import ( cpu_prepack_moe_weight_int8, ) from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( - _CPU_MOE_ACT_FN, - CPUFusedMOE, +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, + biased_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + ArmCPUUnquantizedExperts, + CPUUnquantizedExperts, + X86CPUUnquantizedExperts, + select_experts, ) from vllm.platforms import CpuArchEnum, current_platform from vllm.utils.torch_utils import set_random_seed @@ -36,7 +48,10 @@ ACT = [ ] USE_BIAS = [False, True] ISA = ["vec"] -if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: +if ( + current_platform.get_cpu_architecture() == CpuArchEnum.ARM + and sys.platform != "darwin" +): ISA.append("neon") if torch.cpu._is_amx_tile_supported(): ISA.append("amx") @@ -44,6 +59,27 @@ if torch.cpu._is_amx_tile_supported(): DTYPE = [torch.bfloat16] +def _ref_moe_activation( + input: torch.Tensor, + activation: MoEActivation, +) -> torch.Tensor: + if activation == MoEActivation.SWIGLUOAI: + gate = input[..., ::2].clamp(max=7.0) + up = input[..., 1::2].clamp(min=-7.0, max=7.0) + glu = gate * torch.sigmoid(gate * 1.702) + return (up + 1) * glu + + d = input.shape[-1] // 2 + gate, up = input[..., :d], input[..., d:] + if activation == MoEActivation.SILU: + return torch.nn.functional.silu(gate) * up + if activation == MoEActivation.GELU: + return torch.nn.functional.gelu(gate, approximate="none") * up + if activation == MoEActivation.GELU_TANH: + return torch.nn.functional.gelu(gate, approximate="tanh") * up + raise ValueError(f"Unsupported activation: {activation}") + + def ref_fused_moe( input: torch.Tensor, w13: torch.Tensor, @@ -87,7 +123,8 @@ def ref_fused_moe( tokens_for_this_expert, curr_w13, curr_w13_bias ) # Note: to simulate the kernel implementation - gate_up = _CPU_MOE_ACT_FN[activation](gate_up).to(dtype=input.dtype).float() + gate_up = _ref_moe_activation(gate_up, activation) + gate_up = gate_up.to(dtype=input.dtype).float() expert_out = torch.nn.functional.linear(gate_up, curr_w2, curr_w2_bias) outputs.append(expert_out) @@ -166,7 +203,7 @@ def ref_fused_moe_int8( if w13_bias is not None: gate_up += w13_bias[expert_idx].float() - intermediate = _CPU_MOE_ACT_FN[activation](gate_up).to(input.dtype) + intermediate = _ref_moe_activation(gate_up, activation).to(input.dtype) intermediate_int8, intermediate_scale = quantize_per_token(intermediate) output = torch.matmul( intermediate_int8.float(), @@ -356,8 +393,8 @@ UNALIGNED_INTERMEDIATE_DIM = 176 class _StubMoELayer(torch.nn.Module): """Minimal stand-in for the real MoE layer module, exposing just what - CPUFusedMOE reads/replaces (w13_weight, w2_weight, activation, and - optionally w13_bias/w2_bias).""" + the unquantized CPU experts read/replace (w13_weight, w2_weight, the + router configuration, and optionally w13_bias/w2_bias).""" def __init__( self, @@ -375,12 +412,77 @@ class _StubMoELayer(torch.nn.Module): self.w13_bias = torch.nn.Parameter(w13_bias, requires_grad=False) if w2_bias is not None: self.w2_bias = torch.nn.Parameter(w2_bias, requires_grad=False) + self.use_grouped_topk = False + self.renormalize = False + self.scoring_func = "softmax" + self.custom_routing_function = None + + +def _make_moe_config( + expert_num: int, + hidden_size: int, + intermediate_size: int, + topk_num: int, + dtype: torch.dtype, + act: MoEActivation, + has_bias: bool = False, +) -> FusedMoEConfig: + return FusedMoEConfig( + num_experts=expert_num, + experts_per_token=topk_num, + hidden_dim=hidden_size, + intermediate_size=intermediate_size, + num_local_experts=expert_num, + num_logical_experts=expert_num, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=act, + in_dtype=dtype, + device="cpu", + routing_method=RoutingMethodType.Default, + has_bias=has_bias, + ) + + +@pytest.mark.parametrize( + ("hidden_size", "intermediate_size", "act", "expected"), + [ + (128, 128, MoEActivation.SILU, True), + # Hidden size is not 32-aligned. + (112, 128, MoEActivation.SILU, False), + (128, 176, MoEActivation.SILU, True), + # SwiGLUOAI's interleaved gate/up layout cannot be padded. + (128, 176, MoEActivation.SWIGLUOAI, False), + ], +) +def test_cpu_vec_fused_moe_shape_support( + hidden_size: int, + intermediate_size: int, + act: MoEActivation, + expected: bool, +): + moe_config = _make_moe_config( + expert_num=8, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + topk_num=4, + dtype=torch.bfloat16, + act=act, + ) + supported, _ = CPUUnquantizedExperts.is_supported_config( + CPUUnquantizedExperts, + moe_config, + None, + None, + mk.FusedMoEActivationFormat.Standard, + ) + assert supported is expected @pytest.mark.parametrize("expert_num", EXPERT_NUM) @pytest.mark.parametrize("hidden_size", HIDDEN_DIM) @pytest.mark.parametrize("use_bias", USE_BIAS) @pytest.mark.parametrize("dtype", DTYPE) +@pytest.mark.parametrize("isa", ISA) @pytest.mark.parametrize( "act", [MoEActivation.SILU, MoEActivation.GELU, MoEActivation.GELU_TANH], @@ -391,14 +493,11 @@ def test_cpu_fused_moe_unaligned_intermediate_size( hidden_size: int, use_bias: bool, dtype: torch.dtype, + isa: str, act: MoEActivation, ): - """An unaligned per-partition moe_intermediate_size must still hit the - grouped-gemm fast path via automatic zero-padding, with numerically - correct output, instead of silently falling back to the much slower - per-expert torch loop.""" - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: - pytest.skip("padding is only applied on the x86 AMX/vector kernels") + """CPU kernels handle unaligned intermediate sizes by zero-padding the + weights before prepacking.""" set_random_seed(0) batch_size = 64 @@ -423,9 +522,15 @@ def test_cpu_fused_moe_unaligned_intermediate_size( w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / ( 0.5 * hidden_size**0.5 ) - score = torch.softmax(router_logits, dim=-1, dtype=torch.float32) - topk_weight, topk_ids = torch.topk(score, topk_num) - topk_ids = topk_ids.to(torch.int32) + # Route with the same helper apply() uses internally, so the reference + # only differs from the kernel in how the experts are evaluated. + topk_weight, topk_ids = select_experts( + hidden_states=input, + router_logits=router_logits, + top_k=topk_num, + use_grouped_topk=False, + renormalize=False, + ) ref_output = ref_fused_moe( input, w13, w2, w13_bias, w2_bias, topk_weight, topk_ids, act @@ -439,40 +544,48 @@ def test_cpu_fused_moe_unaligned_intermediate_size( w2_bias.clone() if w2_bias is not None else None, ) - cpu_moe = CPUFusedMOE(layer) - assert cpu_moe.forward_method == cpu_moe.forward_grouped_gemm, ( - "expected the padded intermediate size to hit the grouped-gemm fast path" + moe_config = _make_moe_config( + expert_num, hidden_size, intermediate_size, topk_num, dtype, act, use_bias + ) + experts_cls = { + "vec": CPUUnquantizedExperts, + "amx": X86CPUUnquantizedExperts, + "neon": ArmCPUUnquantizedExperts, + }[isa] + supported, reason = experts_cls.is_supported_config( + experts_cls, + moe_config, + None, + None, + mk.FusedMoEActivationFormat.Standard, + ) + assert supported, ( + f"expected the padded intermediate size to be supported, got {reason}" ) - output = cpu_moe.forward_method( - layer, input, topk_weight, topk_ids, act, expert_num, False + # Mirror UnquantizedFusedMoEMethod._setup_kernel: the quant config is + # built from the layer first, then the experts shuffle the layer into the + # kernel's runtime format. + quant_config = ( + biased_moe_quant_config(layer.w13_bias, layer.w2_bias) + if use_bias + else FusedMoEQuantConfig.make() + ) + experts = experts_cls(moe_config, quant_config) + assert experts.isa == isa + experts.process_weights_after_loading(layer) + + output = experts.apply( + hidden_states=input, + w1=layer.w13_weight, + w2=layer.w2_weight, + router_logits=router_logits, + activation=act, + global_num_experts=expert_num, + expert_map=None, + a1q_scale=None, + apply_router_weight_on_input=False, ) atol, rtol = get_default_atol(output), get_default_rtol(output) torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol) - - -def test_cpu_fused_moe_unaligned_intermediate_size_swigluoai(default_vllm_config): - """swigluoai's interleaved gate/up layout isn't padded automatically. On - AMX-capable CPUs this must raise rather than silently falling back to - the (correct but much slower) per-expert torch loop; elsewhere it should - fall back exactly as before.""" - set_random_seed(0) - expert_num = 8 - hidden_size = 128 - intermediate_size = UNALIGNED_INTERMEDIATE_DIM - dtype = torch.bfloat16 - up_dim = 2 * intermediate_size - - layer = _StubMoELayer( - torch.randn((expert_num, up_dim, hidden_size), dtype=dtype), - torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype), - MoEActivation.SWIGLUOAI, - ) - - if torch.cpu._is_amx_tile_supported(): - with pytest.raises(RuntimeError): - CPUFusedMOE(layer) - else: - cpu_moe = CPUFusedMOE(layer) - assert cpu_moe.forward_method == cpu_moe.forward_torch diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 8e5f6ce79f4..cee8971ae96 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -1482,85 +1482,6 @@ def test_moe_sum_pad_aware(topk: int, dtype: torch.dtype, topk_ids_dtype: torch. opcheck(torch.ops._moe_C.moe_sum, (input, actual, topk_ids, expert_map)) -@pytest.mark.usefixtures("default_vllm_config") -@pytest.mark.parametrize("m", [1, 33]) -@pytest.mark.parametrize("n,k", [(128, 128)]) -@pytest.mark.parametrize("e", [8]) -@pytest.mark.parametrize("topk", [2]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("with_bias", [False, True]) -@pytest.mark.parametrize("activation", [MoEActivation.SILU]) -@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only test") -def test_cpu_fused_moe_basic( - m: int, - n: int, - k: int, - e: int, - topk: int, - dtype: torch.dtype, - with_bias: bool, - activation: MoEActivation, -): - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import CPUFusedMOE - - device = "cpu" - set_random_seed(7) - - a = torch.randn((m, k), device=device, dtype=dtype) / 10 - w13 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / 10 - w2 = torch.randn((e, k, n), device=device, dtype=dtype) / 10 - router_logits = torch.randn((m, e), device=device, dtype=dtype) - - b1 = b2 = None - if with_bias: - b1 = torch.randn((e, 2 * n), device=device, dtype=dtype) / 10 - b2 = torch.randn((e, k), device=device, dtype=dtype) / 10 - - ref = ( - torch_moe(a, w13, w2, router_logits, topk, b1, b2) - if with_bias - else torch_moe(a, w13, w2, router_logits, topk) - ) - - class _Dummy(torch.nn.Module): - def __init__(self, w13, w2, b1=None, b2=None): - super().__init__() - self.w13_weight = torch.nn.Parameter(w13, requires_grad=False) - self.w2_weight = torch.nn.Parameter(w2, requires_grad=False) - if b1 is not None: - self.w13_bias = torch.nn.Parameter(b1, requires_grad=False) - if b2 is not None: - self.w2_bias = torch.nn.Parameter(b2, requires_grad=False) - - layer = _Dummy(w13, w2, b1, b2).to(dtype) - fused = CPUFusedMOE(layer) - out = fused( - layer=layer, - x=a, - use_grouped_topk=False, - top_k=topk, - router_logits=router_logits, - renormalize=False, - global_num_experts=e, - expert_map=None, - custom_routing_function=None, - scoring_func="softmax", - routed_scaling_factor=1.0, - e_score_correction_bias=None, - apply_router_weight_on_input=False, - activation=activation, - ) - - # Tolerances: fp32 tight; bf16 looser (esp. with bias) - if dtype == torch.float32: - atol = 1e-3 - elif with_bias: - atol = 8e-2 - else: - atol = 5e-2 - torch.testing.assert_close(out, ref, atol=atol, rtol=0) - - def _batched_fused_marlin_moe_cases() -> list[Any]: cases = [ pytest.param( diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py index cdff3667ad3..d4c037d978a 100644 --- a/tests/kernels/moe/test_unquantized_backend_selection.py +++ b/tests/kernels/moe/test_unquantized_backend_selection.py @@ -3,14 +3,16 @@ from unittest.mock import patch import pytest +import torch from tests.kernels.moe.utils import make_dummy_moe_config from vllm.model_executor.layers.fused_moe.config import RoutingMethodType from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( UnquantizedMoeBackend, + backend_to_kernel_cls, select_unquantized_moe_backend, ) -from vllm.platforms import current_platform +from vllm.platforms import CpuArchEnum, current_platform skipif_not_cuda_rocm = pytest.mark.skipif( not (current_platform.is_cuda() or current_platform.is_rocm()), @@ -18,6 +20,109 @@ skipif_not_cuda_rocm = pytest.mark.skipif( ) +@pytest.mark.parametrize( + ("amx_supported", "in_dtype", "expect_amx_kernel"), + [ + (True, torch.bfloat16, True), + (False, torch.bfloat16, False), + (True, torch.float16, False), + ], +) +def test_x86_cpu_unquantized_kernel_selection( + amx_supported: bool, + in_dtype: torch.dtype, + expect_amx_kernel: bool, +): + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + CPUUnquantizedExperts, + X86CPUUnquantizedExperts, + ) + + with ( + patch.object(current_platform, "is_cpu", return_value=True), + patch.object( + current_platform, + "get_cpu_architecture", + return_value=CpuArchEnum.X86, + ), + patch("torch.cpu._is_amx_tile_supported", return_value=amx_supported), + ): + moe_config = make_dummy_moe_config( + hidden_dim=128, + intermediate_size=128, + in_dtype=in_dtype, + ) + kernel_cls = next( + cls + for cls in backend_to_kernel_cls(UnquantizedMoeBackend.CPU) + if cls.is_supported_config( + cls, + moe_config, + None, + None, + CPUUnquantizedExperts.activation_format(), + )[0] + ) + + expected_kernel_cls = ( + X86CPUUnquantizedExperts if expect_amx_kernel else CPUUnquantizedExperts + ) + assert kernel_cls is expected_kernel_cls + + +@pytest.mark.parametrize( + ("platform", "in_dtype", "expect_arm_kernel"), + [ + ("linux", torch.bfloat16, True), + ("linux", torch.float16, False), + ("darwin", torch.bfloat16, False), + ], +) +def test_arm_cpu_unquantized_kernel_selection( + platform: str, + in_dtype: torch.dtype, + expect_arm_kernel: bool, +): + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + ArmCPUUnquantizedExperts, + CPUUnquantizedExperts, + ) + + with ( + patch.object(current_platform, "is_cpu", return_value=True), + patch.object( + current_platform, + "get_cpu_architecture", + return_value=CpuArchEnum.ARM, + ), + patch( + "vllm.model_executor.layers.fused_moe.experts.cpu_moe.sys.platform", + platform, + ), + ): + moe_config = make_dummy_moe_config( + hidden_dim=128, + intermediate_size=128, + in_dtype=in_dtype, + ) + kernel_cls = next( + cls + for cls in backend_to_kernel_cls(UnquantizedMoeBackend.CPU) + if cls.is_supported_config( + cls, + moe_config, + None, + None, + CPUUnquantizedExperts.activation_format(), + )[0] + ) + + expected_kernel_cls = ( + ArmCPUUnquantizedExperts if expect_arm_kernel else CPUUnquantizedExperts + ) + assert kernel_cls is expected_kernel_cls + + @pytest.mark.parametrize( "platform_method,expected_backend", [ @@ -69,14 +174,20 @@ def test_select_default_backend_by_platform( patch.object(current_platform, "is_out_of_tree", return_value=False), patch.object(current_platform, platform_method, return_value=True), ): - moe_config = make_dummy_moe_config() + # CPU's grouped-gemm kernels require hidden/intermediate sizes + # aligned to 32; the size-1 defaults only work for backends that + # don't check shapes at selection time. + moe_config = ( + make_dummy_moe_config(hidden_dim=128, intermediate_size=128) + if expected_backend == UnquantizedMoeBackend.CPU + else make_dummy_moe_config() + ) selected_backend, expert_cls = select_unquantized_moe_backend( moe_config=moe_config ) assert selected_backend == expected_backend if expected_backend in [ - UnquantizedMoeBackend.CPU, UnquantizedMoeBackend.OOT, UnquantizedMoeBackend.TPU, ]: diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py deleted file mode 100644 index c89e5ddc304..00000000000 --- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +++ /dev/null @@ -1,577 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import weakref -from collections.abc import Callable - -import torch -from torch.nn import functional as F - -from vllm import _custom_ops as ops -from vllm._custom_ops import ( - CPUQuantMethod, - cpu_fused_moe, - cpu_prepack_moe_weight, - fused_experts_cpu, -) -from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.quantization.utils.layer_utils import replace_parameter -from vllm.platforms import CpuArchEnum, current_platform -from vllm.utils.torch_utils import direct_register_custom_op - -_CPU_MOE_LAYER_CACHE = {} -# The CPU grouped-gemm MoE kernels (AMX and vector) tile the expert -# intermediate ("N") dimension in blocks of this size and have no tail/ -# remainder handling, so a shard is only eligible for the fast path when -# its per-partition intermediate size is a multiple of it. -_MOE_GROUPED_GEMM_N_TILE = 32 - - -def _swigluoai_forward_native( - x: torch.Tensor, - alpha: float = 1.702, - limit: float = 7.0, -) -> torch.Tensor: - """PyTorch-native implementation of SwigluOAIAndMul.forward_native. - - Standalone function to avoid instantiating SwigluOAIAndMul (a CustomOp) - which would trigger get_current_vllm_config() before config is set. - """ - gate, up = x[..., ::2], x[..., 1::2] - gate = gate.clamp(min=None, max=limit) - up = up.clamp(min=-limit, max=limit) - glu = gate * torch.sigmoid(gate * alpha) - gated_output = (up + 1) * glu - return gated_output - - -def _gelu_and_mul( - x: torch.Tensor, -) -> torch.Tensor: - d = x.shape[-1] // 2 - return F.gelu(x[..., :d], approximate="none") * x[..., d:] - - -# Map activation names to their native forward functions. -# Uses static methods or standalone functions to avoid instantiating CustomOp -# classes, which would call get_current_vllm_config() before config is set. -_CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = { - MoEActivation.SILU: SiluAndMul.forward_native, - MoEActivation.SWIGLUOAI: _swigluoai_forward_native, - MoEActivation.GELU: _gelu_and_mul, - MoEActivation.GELU_TANH: ( - lambda x: F.gelu(x[..., : x.shape[-1] // 2], approximate="tanh") - * x[..., x.shape[-1] // 2 :] - ), -} - - -def grouped_topk( - hidden_states: torch.Tensor, - gating_output: torch.Tensor, - topk: int, - renormalize: bool, - num_expert_group: int = 0, - topk_group: int = 0, - scoring_func: str = "softmax", - routed_scaling_factor: float = 1.0, - e_score_correction_bias: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" - - gating_output = gating_output.float() - if scoring_func == "softmax": - scores = torch.softmax(gating_output, dim=-1) - elif scoring_func == "sigmoid": - scores = gating_output.sigmoid() - else: - raise ValueError(f"Unsupported scoring function: {scoring_func}") - - num_token = scores.shape[0] - if e_score_correction_bias is not None: - original_scores = scores - scores = scores + e_score_correction_bias.unsqueeze(0) - group_scores = ( - scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1) - ) - else: - group_scores = ( - scores.view(num_token, num_expert_group, -1).max(dim=-1).values - ) # [n, n_group] - group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[ - 1 - ] # [n, top_k_group] - group_mask = torch.zeros_like(group_scores) # [n, n_group] - group_mask.scatter_(1, group_idx, 1) # [n, n_group] - score_mask = ( - group_mask.unsqueeze(-1) - .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group) - .reshape(num_token, -1) - ) # [n, e] - tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) # [n, e] - - if e_score_correction_bias is not None: - topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1] - topk_weights = original_scores.gather(1, topk_ids) - else: - topk_weights, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False) - - if renormalize: - topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) - - if routed_scaling_factor != 1.0: - topk_weights = topk_weights * routed_scaling_factor - return topk_weights, topk_ids.to(torch.int32) - - -def select_experts( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - top_k: int, - use_grouped_topk: bool, - renormalize: bool, - topk_group: int | None = None, - num_expert_group: int | None = None, - custom_routing_function: Callable | None = None, - scoring_func: str = "softmax", - routed_scaling_factor: float = 1.0, - e_score_correction_bias: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - if use_grouped_topk: - assert topk_group is not None - assert num_expert_group is not None - return grouped_topk( - hidden_states=hidden_states, - gating_output=router_logits, - topk=top_k, - renormalize=renormalize, - num_expert_group=num_expert_group, - topk_group=topk_group, - scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, - e_score_correction_bias=e_score_correction_bias, - ) - elif custom_routing_function is None: - assert scoring_func == "softmax" - topk_logit_vals, topk_idx = torch.topk( - router_logits, k=top_k, dim=-1, sorted=False - ) - if renormalize: - topk_vals = torch.softmax(topk_logit_vals, dim=-1) - else: - logZ = torch.logsumexp(router_logits, dim=-1, keepdim=True) - topk_vals = (topk_logit_vals - logZ).exp() - return topk_vals.to(torch.float32), topk_idx.to(torch.int32) - else: - return custom_routing_function( - hidden_states=hidden_states, - gating_output=router_logits, - topk=top_k, - renormalize=renormalize, - ) - - -class SGLFusedMOE: - def __init__(self, layer: torch.nn.Module) -> None: - pass - - def __call__( - self, - layer: torch.nn.Module, - x: torch.Tensor, - use_grouped_topk: bool, - top_k: int, - router_logits: torch.Tensor, - renormalize: bool, - topk_group: int | None = None, - num_expert_group: int | None = None, - global_num_experts: int = -1, - expert_map: torch.Tensor | None = None, - custom_routing_function: Callable | None = None, - scoring_func: str = "softmax", - routed_scaling_factor: float = 1.0, - e_score_correction_bias: torch.Tensor | None = None, - apply_router_weight_on_input: bool = False, - activation: MoEActivation = MoEActivation.SILU, - ) -> torch.Tensor: - assert activation == MoEActivation.SILU, f"{activation} is not supported." - assert not apply_router_weight_on_input - topk_weights, topk_ids = select_experts( - hidden_states=x, - router_logits=router_logits, - use_grouped_topk=use_grouped_topk, - top_k=top_k, - renormalize=renormalize, - topk_group=topk_group, - num_expert_group=num_expert_group, - custom_routing_function=custom_routing_function, - scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, - e_score_correction_bias=e_score_correction_bias, - ) - - return fused_experts_cpu( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - False, # inplace - CPUQuantMethod.UNQUANT, # moe_comp_method - None, # w1_scale - None, # w2_scale - None, # w1_zero - None, # w2_zero - None, # block_size - None, # w1_bias - None, # w2_bias - None, # alpha - None, # limit - True, # is_vnni - ) - - -class CPUFusedMOE: - """CPU-based fused MoE implementation.""" - - def __init__(self, layer: torch.nn.Module) -> None: - self._pad_moe_intermediate_for_grouped_gemm(layer) - use_grouped_gemm, isa = self.check_grouped_gemm(layer) - self.isa = isa - if use_grouped_gemm: - self.forward_method = self.forward_grouped_gemm - self.init_moe_grouped_gemm(layer=layer) - else: - self.forward_method = self.forward_torch - self.init_moe_torch(layer=layer) - - def __call__( - self, - layer: torch.nn.Module, - x: torch.Tensor, - use_grouped_topk: bool, - top_k: int, - router_logits: torch.Tensor, - renormalize: bool, - topk_group: int | None = None, - num_expert_group: int | None = None, - global_num_experts: int = -1, - expert_map: torch.Tensor | None = None, - custom_routing_function: Callable | None = None, - scoring_func: str = "softmax", - routed_scaling_factor: float = 1.0, - e_score_correction_bias: torch.Tensor | None = None, - apply_router_weight_on_input: bool = False, - activation: MoEActivation = MoEActivation.SILU, - ) -> torch.Tensor: - assert activation in _CPU_MOE_ACT_FN, f"{activation} is not supported." - - topk_weights, topk_ids = select_experts( - hidden_states=x, - router_logits=router_logits, - use_grouped_topk=use_grouped_topk, - top_k=top_k, - renormalize=renormalize, - topk_group=topk_group, - num_expert_group=num_expert_group, - custom_routing_function=custom_routing_function, - scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, - e_score_correction_bias=e_score_correction_bias, - ) - - return self.forward_method( - layer, - x, - topk_weights, - topk_ids, - activation, - global_num_experts, - apply_router_weight_on_input, - ) - - def _pad_moe_intermediate_for_grouped_gemm(self, layer: torch.nn.Module) -> None: - """Zero-pad the per-partition MoE intermediate dim up to a multiple - of _MOE_GROUPED_GEMM_N_TILE, so the AMX/vector grouped-gemm kernels - can be used even when TP sharding (moe_intermediate_size // tp_size) - lands on an unaligned value (e.g. moe_intermediate_size=704 at tp=4 - -> 176). Only applies to the x86 (AMX/vec) kernels and half-split - gate/up activations; interleaved layouts (swigluoai) are left - untouched. - """ - if not hasattr(torch.ops._C, "prepack_moe_weight"): - return - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: - return - - intermediate_size = layer.w2_weight.size(2) - remainder = intermediate_size % _MOE_GROUPED_GEMM_N_TILE - if remainder == 0: - return - if layer.activation == MoEActivation.SWIGLUOAI: - return - - pad = _MOE_GROUPED_GEMM_N_TILE - remainder - padded_size = intermediate_size + pad - num_experts, _, hidden_size = layer.w13_weight.shape - - new_w13 = layer.w13_weight.new_zeros(num_experts, 2 * padded_size, hidden_size) - new_w13[:, :intermediate_size] = layer.w13_weight[:, :intermediate_size] - new_w13[:, padded_size : padded_size + intermediate_size] = layer.w13_weight[ - :, intermediate_size: - ] - replace_parameter(layer, "w13_weight", new_w13) - - new_w2 = layer.w2_weight.new_zeros(num_experts, hidden_size, padded_size) - new_w2[:, :, :intermediate_size] = layer.w2_weight - replace_parameter(layer, "w2_weight", new_w2) - - if hasattr(layer, "w13_bias"): - new_bias = layer.w13_bias.new_zeros(num_experts, 2 * padded_size) - new_bias[:, :intermediate_size] = layer.w13_bias[:, :intermediate_size] - new_bias[:, padded_size : padded_size + intermediate_size] = layer.w13_bias[ - :, intermediate_size: - ] - replace_parameter(layer, "w13_bias", new_bias) - - def _grouped_gemm_alignment_error(self, layer: torch.nn.Module) -> str: - # w2's input size is the per-partition MoE intermediate size (the - # dimension TP-sharding splits), and it's what most commonly breaks - # alignment, e.g. moe_intermediate_size=704 at tp=4 gives - # 704 // 4 == 176, which isn't a multiple of 32. - intermediate_size_per_partition = layer.w2_weight.size(2) - return ( - "CPU fused-MoE AMX grouped-gemm kernel cannot be used for a " - f"layer with w13 shape {tuple(layer.w13_weight.shape)} / w2 " - f"shape {tuple(layer.w2_weight.shape)}: the per-partition MoE " - f"intermediate size ({intermediate_size_per_partition}) is not " - f"a multiple of {_MOE_GROUPED_GEMM_N_TILE}, and automatic " - "zero-padding could not resolve this (typically because the " - "activation uses an interleaved gate/up layout, e.g. " - "swigluoai). vLLM refuses to silently fall back to the much " - "slower per-expert torch loop on AMX-capable CPUs; consider a " - "different --tensor-parallel-size." - ) - - def check_grouped_gemm( - self, - layer: torch.nn.Module, - ) -> tuple[bool, str]: - if not hasattr(torch.ops._C, "prepack_moe_weight"): - return False, "none" - - dtype = layer.w13_weight.dtype - w13_input_size = layer.w13_weight.size(2) - w13_output_size = layer.w13_weight.size(1) - w2_input_size = layer.w2_weight.size(2) - w2_output_size = layer.w2_weight.size(1) - - supports_amx = torch.cpu._is_amx_tile_supported() - if supports_amx: - if ( - dtype == torch.bfloat16 - and w13_output_size % 32 == 0 - and w2_output_size % 32 == 0 - and w13_input_size % 32 == 0 - and w2_input_size % 32 == 0 - ): - return True, "amx" - raise RuntimeError(self._grouped_gemm_alignment_error(layer)) - - if not (w13_output_size % 32 == 0 and w2_output_size % 32 == 0): - return False, "none" - - if ( - layer.activation == MoEActivation.SWIGLUOAI - and w2_input_size % _MOE_GROUPED_GEMM_N_TILE != 0 - ): - return False, "none" - - supports_neon = current_platform.get_cpu_architecture() == CpuArchEnum.ARM - if supports_neon: - if ( - dtype == torch.bfloat16 - and w13_input_size % 4 == 0 - and w2_input_size % 4 == 0 - ): - return True, "neon" - return False, "none" - - return True, "vec" - - def init_moe_grouped_gemm( - self, - layer: torch.nn.Module, - ) -> None: - new_w13 = cpu_prepack_moe_weight(layer.w13_weight, self.isa) - replace_parameter(layer, "w13_weight", new_w13) - new_w2 = cpu_prepack_moe_weight(layer.w2_weight, self.isa) - replace_parameter(layer, "w2_weight", new_w2) - - def init_moe_torch( - self, - layer: torch.nn.Module, - ) -> None: - use_onednn_mm = ops._supports_onednn and ops.is_onednn_acl_supported() - num_experts = layer.w13_weight.size(0) - has_w13_bias = hasattr(layer, "w13_bias") - has_w2_bias = hasattr(layer, "w2_bias") - - layer.gate_up_linear = [] - layer.down_linear = [] - - for i in range(num_experts): - layer_w13_weight = layer.w13_weight[i] - layer_w13_bias = layer.w13_bias[i] if has_w13_bias else None - layer_w2_weight = layer.w2_weight[i] - layer_w2_bias = layer.w2_bias[i] if has_w2_bias else None - if use_onednn_mm: - gate_up_handle = ops.create_onednn_mm(layer_w13_weight.t(), 32) - layer.gate_up_linear.append( - lambda x, handle=gate_up_handle, bias=layer_w13_bias: ops.onednn_mm( - handle, x, bias - ) - ) - down_handle = ops.create_onednn_mm(layer_w2_weight.t(), 32) - layer.down_linear.append( - lambda x, handle=down_handle, bias=layer_w2_bias: ops.onednn_mm( - handle, x, bias - ) - ) - else: - layer.gate_up_linear.append( - lambda x, w=layer_w13_weight, b=layer_w13_bias: F.linear(x, w, b) - ) - layer.down_linear.append( - lambda x, w=layer_w2_weight, b=layer_w2_bias: F.linear(x, w, b) - ) - - if use_onednn_mm: # remove weight - layer.w13_weight = torch.nn.Parameter(torch.empty(0), requires_grad=False) - layer.w2_weight = torch.nn.Parameter(torch.empty(0), requires_grad=False) - - _CPU_MOE_LAYER_CACHE[id(layer)] = weakref.ref(layer) - - def forward_grouped_gemm( - self, - layer: torch.nn.Module, - input: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - activation: MoEActivation, - global_num_experts: int = -1, - skip_weighted: bool = False, - ) -> torch.Tensor: - if skip_weighted: - assert topk_ids.size(1) == 1, ( - "apply_router_weight_on_input is only implemented for topk=1" - ) - input.mul_(topk_weights.to(input.dtype)) - - output = cpu_fused_moe( - input, - layer.w13_weight, - layer.w2_weight, - getattr(layer, "w13_bias", None), - getattr(layer, "w2_bias", None), - topk_weights, - topk_ids, - activation.value, - self.isa, - skip_weighted, - ) - return output - - def forward_torch( - self, - layer: torch.nn.Module, - input: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - activation: MoEActivation, - global_num_experts: int = -1, - skip_weighted: bool = False, - ) -> torch.Tensor: - if skip_weighted: - assert topk_ids.size(1) == 1, ( - "apply_router_weight_on_input is only implemented for topk=1" - ) - input.mul_(topk_weights.to(input.dtype)) - - output = torch.empty_like(input) - layer_id = id(layer) - torch.ops.vllm.cpu_fused_moe_torch( - layer_id, - output, - input, - topk_weights, - topk_ids, - activation.value, - global_num_experts, - skip_weighted, - ) - - return output - - -def cpu_fused_moe_torch( - layer_id: int, - output: torch.Tensor, - input: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - activation: str, - global_num_experts: int = -1, - skip_weighted: bool = False, -) -> None: - act = MoEActivation.from_str(activation) - layer = _CPU_MOE_LAYER_CACHE[layer_id]() - - # Ref code from https://github.com/sgl-project/sglang/blob/716e682721397df103f347d22da8bd46c6016dab/python/sglang/srt/layers/moe/fused_moe_native.py#L53 - len_experts = global_num_experts - - cnts = topk_ids.new_zeros((topk_ids.shape[0], len_experts)) - cnts.scatter_(1, topk_ids.to(torch.int64), 1) - tokens_per_expert = cnts.sum(dim=0) - idxs = topk_ids.view(-1).argsort() - - sorted_tokens = input[idxs // topk_ids.shape[1]] - tokens_per_expert = tokens_per_expert.cpu().numpy() - - outputs = [] - start_idx = 0 - - for i, num_tokens in enumerate(tokens_per_expert): - end_idx = start_idx + num_tokens - if num_tokens == 0: - continue - tokens_for_this_expert = sorted_tokens[start_idx:end_idx] - - gate_up = layer.gate_up_linear[i](tokens_for_this_expert) # type: ignore - gate_up = _CPU_MOE_ACT_FN[act](gate_up) - expert_out = layer.down_linear[i](gate_up) # type: ignore - outputs.append(expert_out) - start_idx = end_idx - - outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) - new_x = torch.empty_like(outs) - - new_x[idxs] = outs - if skip_weighted: - final_out = new_x - else: - final_out = ( - new_x.view(*topk_ids.shape, -1) - .type(topk_weights.dtype) - .mul_(topk_weights.unsqueeze(dim=-1)) - .sum(dim=1) - .type(new_x.dtype) - ) - output.copy_(final_out) - - -direct_register_custom_op( - op_name="cpu_fused_moe_torch", - op_func=cpu_fused_moe_torch, - mutates_args=["output"], -) diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py index f4933e1bd76..5816bc5d3a5 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py @@ -218,7 +218,7 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): Returns: Output tensor after MoE computation """ - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( select_experts, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index 03561ceb8a2..ab8a3a3b7a6 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -1,6 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""CPU quantized fused MoE experts.""" +"""CPU fused MoE experts.""" + +import math +import sys +from collections.abc import Callable +from typing import cast import torch @@ -9,7 +14,9 @@ from vllm._custom_ops import ( CPUQuantAlgo, CPUQuantMethod, convert_weight_packed_scale_zp, + cpu_fused_moe, cpu_fused_moe_int8, + cpu_prepack_moe_weight, cpu_prepack_moe_weight_int8, fused_experts_cpu, ) @@ -29,7 +36,439 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kInt8StaticChannelSym, kMxfp4Static, ) +from vllm.model_executor.utils import replace_parameter from vllm.platforms import CpuArchEnum, current_platform +from vllm.utils.math_utils import round_up + +# =========================================================================== +# Routing +# =========================================================================== + + +def grouped_topk( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + num_expert_group: int = 0, + topk_group: int = 0, + scoring_func: str = "softmax", + routed_scaling_factor: float = 1.0, + e_score_correction_bias: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" + + gating_output = gating_output.float() + if scoring_func == "softmax": + scores = torch.softmax(gating_output, dim=-1) + elif scoring_func == "sigmoid": + scores = gating_output.sigmoid() + else: + raise ValueError(f"Unsupported scoring function: {scoring_func}") + + num_token = scores.shape[0] + if e_score_correction_bias is not None: + original_scores = scores + scores = scores + e_score_correction_bias.unsqueeze(0) + group_scores = ( + scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1) + ) + else: + group_scores = ( + scores.view(num_token, num_expert_group, -1).max(dim=-1).values + ) # [n, n_group] + group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[ + 1 + ] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, n_group] + group_mask.scatter_(1, group_idx, 1) # [n, n_group] + score_mask = ( + group_mask.unsqueeze(-1) + .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group) + .reshape(num_token, -1) + ) # [n, e] + tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) # [n, e] + + if e_score_correction_bias is not None: + topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1] + topk_weights = original_scores.gather(1, topk_ids) + else: + topk_weights, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False) + + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + + if routed_scaling_factor != 1.0: + topk_weights = topk_weights * routed_scaling_factor + return topk_weights, topk_ids.to(torch.int32) + + +def select_experts( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + top_k: int, + use_grouped_topk: bool, + renormalize: bool, + topk_group: int | None = None, + num_expert_group: int | None = None, + custom_routing_function: Callable | None = None, + scoring_func: str = "softmax", + routed_scaling_factor: float = 1.0, + e_score_correction_bias: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + if use_grouped_topk: + assert topk_group is not None + assert num_expert_group is not None + return grouped_topk( + hidden_states=hidden_states, + gating_output=router_logits, + topk=top_k, + renormalize=renormalize, + num_expert_group=num_expert_group, + topk_group=topk_group, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + e_score_correction_bias=e_score_correction_bias, + ) + elif custom_routing_function is None: + assert scoring_func == "softmax" + topk_logit_vals, topk_idx = torch.topk( + router_logits, k=top_k, dim=-1, sorted=False + ) + if renormalize: + topk_vals = torch.softmax(topk_logit_vals, dim=-1) + else: + logZ = torch.logsumexp(router_logits, dim=-1, keepdim=True) + topk_vals = (topk_logit_vals - logZ).exp() + return topk_vals.to(torch.float32), topk_idx.to(torch.int32) + else: + topk_weights, topk_ids = custom_routing_function( + hidden_states=hidden_states, + gating_output=router_logits, + topk=top_k, + renormalize=renormalize, + ) + # cpu_fused_moe reads routing tensors as contiguous float32/int32 + # buffers and does not account for tensor strides. + topk_weights = topk_weights.to(torch.float32).contiguous() + topk_ids = topk_ids.to(torch.int32).contiguous() + return topk_weights, topk_ids + + +# =========================================================================== +# Unquantized (BF16/FP16/FP32) MoE +# =========================================================================== + + +class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic): + """Portable vector grouped-gemm unquantized MoE experts.""" + + isa = "vec" + output_alignment = 32 + reduction_alignment = 1 + + @classmethod + def _intermediate_alignment(cls) -> int: + return math.lcm(cls.output_alignment, cls.reduction_alignment) + + @classmethod + def _padded_intermediate_size(cls, moe_config: FusedMoEConfig) -> int: + intermediate_size = moe_config.intermediate_size_per_partition + if moe_config.activation == MoEActivation.SWIGLUOAI: + return intermediate_size + return round_up(intermediate_size, cls._intermediate_alignment()) + + @classmethod + def _supports_grouped_gemm( + cls, + moe_config: FusedMoEConfig, + ) -> tuple[bool, str | None]: + intermediate_size = cls._padded_intermediate_size(moe_config) + if ( + moe_config.hidden_dim % cls.output_alignment != 0 + or intermediate_size % cls.output_alignment != 0 + ): + return False, ( + "kernel requires hidden and intermediate dimensions divisible by " + f"{cls.output_alignment}" + ) + if ( + moe_config.hidden_dim % cls.reduction_alignment != 0 + or intermediate_size % cls.reduction_alignment != 0 + ): + return False, ( + "kernel requires reduction dimensions divisible by " + f"{cls.reduction_alignment}" + ) + return True, None + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + # Router configuration that the monolithic apply() signature cannot + # carry. Captured off the layer in process_weights_after_loading. + self.use_grouped_topk = False + self.renormalize = False + self.scoring_func = "softmax" + self.custom_routing_function: Callable | None = None + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + return current_platform.is_cpu() + + @staticmethod + def is_supported_config( + cls: type[mk.FusedMoEExperts], + moe_config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: mk.FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + supported, reason = mk.FusedMoEExperts.is_supported_config( + cls, moe_config, weight_key, activation_key, activation_format + ) + if not supported: + return supported, reason + cpu_cls = cast(type[CPUUnquantizedExperts], cls) + return cpu_cls._supports_grouped_gemm(moe_config) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in ( + MoEActivation.SILU, + MoEActivation.SWIGLUOAI, + MoEActivation.GELU, + MoEActivation.GELU_TANH, + ) + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == (None, None) + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return True + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + # Routing runs in select_experts(), which covers every routing method + # a layer can be configured with, including custom routing functions. + return True + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return True + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self._pad_moe_intermediate(layer) + self.use_grouped_topk = layer.use_grouped_topk + self.renormalize = layer.renormalize + self.scoring_func = layer.scoring_func + self.custom_routing_function = layer.custom_routing_function + replace_parameter( + layer, "w13_weight", cpu_prepack_moe_weight(layer.w13_weight, self.isa) + ) + replace_parameter( + layer, "w2_weight", cpu_prepack_moe_weight(layer.w2_weight, self.isa) + ) + + def _pad_moe_intermediate(self, layer: torch.nn.Module) -> None: + """Zero-pad the per-partition MoE intermediate dim of both weights and + the expert bias, see `_padded_intermediate_size`.""" + intermediate_size = self.moe_config.intermediate_size_per_partition + padded_size = self._padded_intermediate_size(self.moe_config) + if padded_size == intermediate_size: + return + + num_experts, _, hidden_size = layer.w13_weight.shape + + new_w13 = layer.w13_weight.new_zeros(num_experts, 2 * padded_size, hidden_size) + new_w13[:, :intermediate_size] = layer.w13_weight[:, :intermediate_size] + new_w13[:, padded_size : padded_size + intermediate_size] = layer.w13_weight[ + :, intermediate_size: + ] + replace_parameter(layer, "w13_weight", new_w13) + + new_w2 = layer.w2_weight.new_zeros(num_experts, hidden_size, padded_size) + new_w2[:, :, :intermediate_size] = layer.w2_weight + replace_parameter(layer, "w2_weight", new_w2) + + if hasattr(layer, "w13_bias"): + new_bias = layer.w13_bias.new_zeros(num_experts, 2 * padded_size) + new_bias[:, :intermediate_size] = layer.w13_bias[:, :intermediate_size] + new_bias[:, padded_size : padded_size + intermediate_size] = layer.w13_bias[ + :, intermediate_size: + ] + # Assign through .data rather than replacing the Parameter: the + # quant config is built before this runs and holds a reference to + # this very object, which is what feeds self.w1_bias in apply(). + layer.w13_bias.data = new_bias + + def _select_experts( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + num_expert_group: int | None, + topk_group: int | None, + e_score_correction_bias: torch.Tensor | None, + routed_scaling_factor: float | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + top_k=self.moe_config.experts_per_token, + use_grouped_topk=self.use_grouped_topk, + renormalize=self.renormalize, + topk_group=topk_group, + num_expert_group=num_expert_group, + custom_routing_function=self.custom_routing_function, + scoring_func=self.scoring_func, + routed_scaling_factor=( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ), + e_score_correction_bias=e_score_correction_bias, + ) + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + topk_weights, topk_ids = self._select_experts( + hidden_states, + router_logits, + num_expert_group, + topk_group, + e_score_correction_bias, + routed_scaling_factor, + ) + + if apply_router_weight_on_input: + assert topk_ids.size(1) == 1, ( + "apply_router_weight_on_input is only implemented for topk=1" + ) + hidden_states.mul_(topk_weights.to(hidden_states.dtype)) + + return cpu_fused_moe( + hidden_states, + w1, + w2, + self.w1_bias, + self.w2_bias, + topk_weights, + topk_ids, + activation.value, + self.isa, + apply_router_weight_on_input, + ) + + +class X86CPUUnquantizedExperts(CPUUnquantizedExperts): + """x86 AMX grouped-gemm unquantized MoE experts.""" + + isa = "amx" + output_alignment = 32 + reduction_alignment = 32 + + @staticmethod + def _supports_current_device() -> bool: + return ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.X86 + and torch.cpu._is_amx_tile_supported() + ) + + @staticmethod + def is_supported_config( + cls: type[mk.FusedMoEExperts], + moe_config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: mk.FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + supported, reason = mk.FusedMoEExperts.is_supported_config( + cls, moe_config, weight_key, activation_key, activation_format + ) + if not supported: + return supported, reason + if moe_config.in_dtype != torch.bfloat16: + return False, "kernel requires bfloat16 activations" + cpu_cls = cast(type[CPUUnquantizedExperts], cls) + return cpu_cls._supports_grouped_gemm(moe_config) + + +class ArmCPUUnquantizedExperts(CPUUnquantizedExperts): + """Arm NEON grouped-gemm unquantized MoE experts.""" + + isa = "neon" + output_alignment = 32 + reduction_alignment = 4 + + @staticmethod + def _supports_current_device() -> bool: + return ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.ARM + and sys.platform != "darwin" + ) + + @staticmethod + def is_supported_config( + cls: type[mk.FusedMoEExperts], + moe_config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: mk.FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + supported, reason = mk.FusedMoEExperts.is_supported_config( + cls, moe_config, weight_key, activation_key, activation_format + ) + if not supported: + return supported, reason + if moe_config.in_dtype != torch.bfloat16: + return False, "kernel requires bfloat16 activations" + cpu_cls = cast(type[CPUUnquantizedExperts], cls) + return cpu_cls._supports_grouped_gemm(moe_config) + # =========================================================================== # FP8 W8A16 MoE @@ -131,10 +570,6 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic): routed_scaling_factor: float | None = None, topk_group: int | None = None, ) -> torch.Tensor: - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( - select_experts, - ) - topk_weights, topk_ids = select_experts( hidden_states=hidden_states, router_logits=router_logits, @@ -289,10 +724,6 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic): routed_scaling_factor: float | None = None, topk_group: int | None = None, ) -> torch.Tensor: - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( - select_experts, - ) - topk_weights, topk_ids = select_experts( hidden_states=hidden_states, router_logits=router_logits, @@ -505,10 +936,6 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): "apply_router_weight_on_input=True. " ) - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( - select_experts, - ) - topk_weights, topk_ids = select_experts( hidden_states=hidden_states, router_logits=router_logits, @@ -631,7 +1058,6 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: """VNNI-prepack INT8 MoE weights for CPU kernel.""" - from vllm.model_executor.utils import replace_parameter w13 = torch.ops._C.convert_weight_packed(layer.w13_weight) w2 = torch.ops._C.convert_weight_packed(layer.w2_weight) @@ -655,10 +1081,6 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic): routed_scaling_factor: float | None = None, topk_group: int | None = None, ) -> torch.Tensor: - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( - select_experts, - ) - topk_weights, topk_ids = select_experts( hidden_states=hidden_states, router_logits=router_logits, @@ -796,8 +1218,6 @@ class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic): return True def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - from vllm.model_executor.utils import replace_parameter - w13 = cpu_prepack_moe_weight_int8(layer.w13_weight, "neon") w2 = cpu_prepack_moe_weight_int8(layer.w2_weight, "neon") replace_parameter(layer, "w13_weight", w13) @@ -819,10 +1239,6 @@ class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic): routed_scaling_factor: float | None = None, topk_group: int | None = None, ) -> torch.Tensor: - from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( - select_experts, - ) - topk_weights, topk_ids = select_experts( hidden_states=hidden_states, router_logits=router_logits, diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index bf77a316bb4..6fbcd67279f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -143,6 +143,20 @@ def backend_to_kernel_cls( return [XPUExperts] + elif backend == UnquantizedMoeBackend.CPU: + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + ArmCPUUnquantizedExperts, + CPUUnquantizedExperts, + X86CPUUnquantizedExperts, + ) + + # Prefer architecture-specific kernels before the portable vector path. + return [ + X86CPUUnquantizedExperts, + ArmCPUUnquantizedExperts, + CPUUnquantizedExperts, + ] + else: raise ValueError(f"Unknown unquantized MoE backend: {backend.value}") @@ -199,10 +213,6 @@ def select_unquantized_moe_backend( Note: Shape-specific fallbacks may still occur at runtime. """ - if current_platform.is_cpu(): - # TODO: migrate to MK structure. - return UnquantizedMoeBackend.CPU, None - if current_platform.is_tpu(): return UnquantizedMoeBackend.TPU, None diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index d913e9ad676..2f17e147a6d 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable from typing import TYPE_CHECKING import torch @@ -33,7 +32,6 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( ) from vllm.model_executor.utils import replace_parameter, set_weight_attrs from vllm.platforms import current_platform -from vllm.platforms.interface import CpuArchEnum if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts @@ -54,13 +52,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): moe_config=self.moe, ) - @property - def is_monolithic(self) -> bool: - # Escape hatch for CPU, which stays on the old monolithic path. - if self.unquantized_backend == UnquantizedMoeBackend.CPU: - return True - return super().is_monolithic - @property def supports_eplb(self) -> bool: return True @@ -71,8 +62,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): ): raise ValueError( f"{self.__class__.__name__} uses the new modular kernel initialization " - "logic for all but the CPU backend. CPU backend is monolithic. " - "So this function should not be called." + "logic. So this function should not be called." ) def select_gemm_impl( @@ -201,6 +191,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): routing_tables=layer._expert_routing_tables(), ) + if self.unquantized_backend == UnquantizedMoeBackend.CPU: + # The CPU experts need the layer itself for the setup that + # convert_to_unquantized_kernel_format cannot express, since + # it only sees the two weight tensors: padding and prepacking + # into the grouped-gemm layout (bias included), and capturing + # the router config that monolithic apply() cannot carry. + self.moe_kernel.fused_experts.process_weights_after_loading(layer) + def process_weights_after_loading(self, layer: "RoutedExperts") -> None: super().process_weights_after_loading(layer) @@ -221,38 +219,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): # OOT handles internally. return - elif self.unquantized_backend == UnquantizedMoeBackend.CPU: - # CPU stays on the old path — no oracle, no moe_kernel. - from vllm.model_executor.layers.fused_moe import cpu_fused_moe - - if current_platform.get_cpu_architecture() == CpuArchEnum.X86: - from vllm.model_executor.layers.utils import check_cpu_sgl_kernel - - dtype_w13 = layer.w13_weight.dtype - _, n_w13, k_w13 = layer.w13_weight.size() - dtype_w2 = layer.w2_weight.dtype - _, n_w2, k_w2 = layer.w2_weight.size() - if ( - envs.VLLM_CPU_SGL_KERNEL - and check_cpu_sgl_kernel(n_w13, k_w13, dtype_w13) - and check_cpu_sgl_kernel(n_w2, k_w2, dtype_w2) - ): - packed_w13_weight = torch.ops._C.convert_weight_packed( - layer.w13_weight - ) - assert packed_w13_weight.size() == layer.w13_weight.size() - layer.w13_weight.copy_(packed_w13_weight) - del packed_w13_weight - packed_w2_weight = torch.ops._C.convert_weight_packed( - layer.w2_weight - ) - assert packed_w2_weight.size() == layer.w2_weight.size() - layer.w2_weight.copy_(packed_w2_weight) - self.cpu_fused_moe: Callable = cpu_fused_moe.SGLFusedMOE(layer) - else: - self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer) - else: - self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer) elif self.unquantized_backend == UnquantizedMoeBackend.XPU: w13 = layer.w13_weight w2 = layer.w2_weight @@ -363,39 +329,18 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): input_ids: torch.Tensor | None = None, ) -> torch.Tensor: assert self.is_monolithic - if self.unquantized_backend == UnquantizedMoeBackend.CPU: - assert self.moe_kernel is None - return self.cpu_fused_moe( - layer, - x, - layer.use_grouped_topk, - layer.top_k, - router_logits, - layer.renormalize, - layer.topk_group, - layer.num_expert_group, - layer.global_num_experts, - layer.expert_map, - layer.custom_routing_function, - layer.scoring_func, - layer.routed_scaling_factor, - layer.e_score_correction_bias, - layer.apply_router_weight_on_input, - layer.activation, - ) - else: - assert self.moe_kernel is not None - return self.moe_kernel.apply_monolithic( - x, - layer.w13_weight, - layer.w2_weight, - router_logits, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - num_expert_group=layer.num_expert_group, - topk_group=layer.topk_group, - e_score_correction_bias=layer.e_score_correction_bias, - routed_scaling_factor=layer.routed_scaling_factor, - ) + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, + )