mirror of
https://github.com/vllm-project/vllm.git
synced 2026-08-04 21:08:06 +00:00
[CPU] Migrate unquantized MoE to the modular-kernel experts structure (#50133)
Signed-off-by: jiang1.li <[email protected]> Signed-off-by: Fadi Arafeh <[email protected]> Co-authored-by: Claude Opus 5 (1M context) <[email protected]> Co-authored-by: Fadi Arafeh <[email protected]>
This commit is contained in:
Li, Jiang
co-authored by
Claude Opus 5
Fadi Arafeh
Claude Opus 5
Fadi Arafeh
parent
b3f97dae24
commit
0a6446005d
@@ -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"
|
||||
|
||||
@@ -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<scalar_t>::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
|
||||
|
||||
@@ -678,6 +678,9 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
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<INT8Vec16> {
|
||||
}
|
||||
};
|
||||
|
||||
// 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
|
||||
// ============================================================================
|
||||
|
||||
@@ -333,6 +333,13 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
return FP32Vec16(ret);
|
||||
}
|
||||
|
||||
FP32Vec16 operator-() const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, this](int i) { ret.val[i] = -reg.val[i]; });
|
||||
return FP32Vec16(ret);
|
||||
}
|
||||
|
||||
FP32Vec16 operator/(const FP32Vec16& b) const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
@@ -356,6 +363,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
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<int, VEC_ELEM_NUM>(
|
||||
@@ -363,6 +374,20 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
return FP32Vec16(ret);
|
||||
}
|
||||
|
||||
FP32Vec16 exp() const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, this](int i) { ret.val[i] = std::exp(reg.val[i]); });
|
||||
return FP32Vec16(ret);
|
||||
}
|
||||
|
||||
FP32Vec16 er() const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, this](int i) { ret.val[i] = std::erf(reg.val[i]); });
|
||||
return FP32Vec16(ret);
|
||||
}
|
||||
|
||||
FP32Vec16 tanh() const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
@@ -404,6 +429,25 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
void save(void* ptr) const { *reinterpret_cast<f32x16_t*>(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 <typename T>
|
||||
struct VecType {
|
||||
using vec_type = void;
|
||||
|
||||
@@ -655,6 +655,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
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<FP32Vec16> {
|
||||
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]);
|
||||
|
||||
@@ -6,8 +6,10 @@
|
||||
#include <bit>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
#include <limits>
|
||||
#include <torch/all.h>
|
||||
|
||||
namespace vec_op {
|
||||
|
||||
struct fp8_e4m3_tag {};
|
||||
@@ -727,6 +729,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
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<FP32Vec16> {
|
||||
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<FP32Vec16> {
|
||||
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<INT8Vec16> {
|
||||
}
|
||||
};
|
||||
|
||||
// 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 <typename T>
|
||||
struct VecType {
|
||||
using vec_type = void;
|
||||
@@ -1325,4 +1389,4 @@ FORCE_INLINE void prefetch(const void* addr) {
|
||||
|
||||
}; // namespace vec_op
|
||||
|
||||
#endif
|
||||
#endif
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
#ifndef CPU_TYPES_X86_HPP
|
||||
#define CPU_TYPES_X86_HPP
|
||||
|
||||
#include <cstring>
|
||||
#include <immintrin.h>
|
||||
#include <sleef.h>
|
||||
#include <torch/all.h>
|
||||
@@ -595,6 +596,8 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
|
||||
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<FP32Vec16> {
|
||||
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<INT8Vec16> {
|
||||
for (int i = 0; i < elem_num; ++i) ptr[i] = ar.values[i];
|
||||
}
|
||||
};
|
||||
|
||||
struct INT8Vec64 : public Vec<INT8Vec64> {
|
||||
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<const __m256i*>(ptr))),
|
||||
reg_high(
|
||||
_mm256_loadu_si256(reinterpret_cast<const __m256i*>(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 <typename T>
|
||||
|
||||
@@ -71,8 +71,7 @@ class TileGemm82 {
|
||||
scalar_t* __restrict__ curr_m_a = curr_a;
|
||||
vec_op::unroll_loop<int32_t, M>([&](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<float>(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;
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -40,12 +40,10 @@ struct VecTypeTrait<c10::BFloat16> {
|
||||
using vec_t = vec_op::BF16Vec16;
|
||||
};
|
||||
|
||||
#if !defined(__powerpc__)
|
||||
template <>
|
||||
struct VecTypeTrait<c10::Half> {
|
||||
using vec_t = vec_op::FP16Vec16;
|
||||
};
|
||||
#endif
|
||||
|
||||
struct Counter {
|
||||
std::atomic<int64_t> counter;
|
||||
|
||||
@@ -91,7 +91,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
|
||||
| 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] |
|
||||
| hpc | standard | fp8 | G(128),T | silu | Y | Y | [`HPCExperts`][vllm.model_executor.layers.fused_moe.hpc_moe.HPCExperts] |
|
||||
| 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` |
|
||||
| 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,</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] |
|
||||
| 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] |
|
||||
|
||||
!!! info "Table key"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
]:
|
||||
|
||||
@@ -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"],
|
||||
)
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user