diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index ff2911dc11..6359a96887 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -522,6 +522,12 @@ if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND) ${VLLM_EXT_SRC}) endif() +if (POWER10_FOUND OR POWER11_FOUND) + set(VLLM_EXT_SRC + "csrc/cpu/cpu_fused_moe_int8.cpp" + ${VLLM_EXT_SRC}) +endif() + if(USE_ONEDNN) set(VLLM_EXT_SRC "csrc/cpu/dnnl_kernels.cpp" diff --git a/csrc/cpu/cpu_fused_moe_int8.cpp b/csrc/cpu/cpu_fused_moe_int8.cpp index 649317b478..27d0f6a7d5 100644 --- a/csrc/cpu/cpu_fused_moe_int8.cpp +++ b/csrc/cpu/cpu_fused_moe_int8.cpp @@ -26,10 +26,23 @@ #define NEON_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::NEON: #endif +#if defined(__powerpc64__) + #include "cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp" + #define VSX_DISPATCH(SCALAR_TYPE, ...) \ + case cpu_utils::ISA::VSX: { \ + using gemm_t = \ + cpu_micro_gemm::MicroGemmINT8; \ + return __VA_ARGS__(); \ + } +#else + #define VSX_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::VSX: +#endif + #define CPU_INT8_ISA_DISPATCH_IMPL(ISA_TYPE, SCALAR_TYPE, ...) \ [&] { \ switch (ISA_TYPE) { \ NEON_DISPATCH(SCALAR_TYPE, __VA_ARGS__) \ + VSX_DISPATCH(SCALAR_TYPE, __VA_ARGS__) \ default: { \ TORCH_CHECK(false, "Invalid CPU ISA type."); \ } \ @@ -81,9 +94,9 @@ void fused_moe_int8_impl( const int32_t thread_num = cpu_utils::get_max_threads(); const int32_t w13_input_buffer_size = cpu_utils::round_up<64>( - gemm_m_tile_size * input_size_13 * sizeof(int8_t)); - const int32_t w2_input_buffer_size = - cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 * sizeof(int8_t)); + (gemm_m_tile_size * input_size_13 + 64) * sizeof(int8_t)); + const int32_t w2_input_buffer_size = cpu_utils::round_up<64>( + (gemm_m_tile_size * input_size_2 + 64) * sizeof(int8_t)); const int32_t w13_n_tile_size = [&]() { const int64_t cache_size = cpu_utils::get_available_l2_size(); diff --git a/csrc/cpu/cpu_types_vsx.hpp b/csrc/cpu/cpu_types_vsx.hpp index d46906130e..eaea921d98 100644 --- a/csrc/cpu/cpu_types_vsx.hpp +++ b/csrc/cpu/cpu_types_vsx.hpp @@ -13,10 +13,10 @@ namespace vec_op { struct fp8_e4m3_tag {}; struct fp8_e5m2_tag {}; -#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +// Note: FP16 (Half) is not supported on POWER VSX — no FP16Vec16 type. #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) @@ -416,6 +416,7 @@ struct FP32Vec8 : public Vec { tmp.val[0] = two_x; tmp.val[1] = two_x; FP32Vec8 temp_vec(tmp); + FP32Vec8 exp_vec = temp_vec.exp(); vector float e = temp_vec.exp().reg.val[0]; vector float num = vec_sub(e, one); @@ -634,6 +635,16 @@ struct FP32Vec16 : public Vec { reg.val[3] = vec_ctf(v.reg.val[3], 0); } + explicit FP32Vec16(const float* __restrict__ base, const INT32Vec16& index) { + INT32Vec16::AliasReg idx; + idx.reg = index.reg; + AliasReg ar; + for (int i = 0; i < VEC_ELEM_NUM; ++i) { + ar.values[i] = base[idx.values[i]]; + } + reg = ar.reg; + } + FP32Vec16 operator*(const FP32Vec16& b) const { return FP32Vec16(f32x4x4_t({vec_mul(reg.val[0], b.reg.val[0]), vec_mul(reg.val[1], b.reg.val[1]), @@ -648,6 +659,13 @@ struct FP32Vec16 : public Vec { vec_add(reg.val[3], b.reg.val[3])})); } + FP32Vec16 operator-() const { + const __vector float zero = vec_splats(0.0f); + return FP32Vec16( + f32x4x4_t({vec_sub(zero, reg.val[0]), vec_sub(zero, reg.val[1]), + vec_sub(zero, reg.val[2]), vec_sub(zero, reg.val[3])})); + } + FP32Vec16 operator-(const FP32Vec16& b) const { return FP32Vec16(f32x4x4_t({vec_sub(reg.val[0], b.reg.val[0]), vec_sub(reg.val[1], b.reg.val[1]), @@ -655,11 +673,6 @@ 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]), diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp new file mode 100644 index 0000000000..6033353728 --- /dev/null +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp @@ -0,0 +1,287 @@ +#ifndef CPU_MICRO_GEMM_INT8_VSX_HPP +#define CPU_MICRO_GEMM_INT8_VSX_HPP + +#include +#include +#include +#include +#include + +#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp" + +#include +#include +#include + +namespace cpu_micro_gemm { + +namespace vsx_int8 { + +constexpr int32_t K = 4; +constexpr int32_t Mr = 4; +constexpr int32_t Nr = 32; + +// --------------------------------------------------------------------------- +// 8-Accumulator MMA 4x32 micro-GEMM (POWER10 xvi8ger4pp). +// +// Computes 4 rows x 32 columns (128 outputs) using all 8 hardware accumulators. +// Operates on pre-biased uint8 weights (B + 128) and subtracts 128 * a_sums. +// Works for all m in [1, 4] with zero scalar loop overhead. +// --------------------------------------------------------------------------- +__attribute__((target("cpu=power10"))) FORCE_INLINE void gemm_mma_4x32( + const int8_t* __restrict__ a_packed, const uint8_t* __restrict__ b_panel, + const int32_t* __restrict__ a_sums, int32_t* __restrict__ c_ptr, + const int32_t m, const int32_t k_size, const int64_t ldc) { + TORCH_CHECK(m > 0 && m <= Mr); + TORCH_CHECK(k_size % K == 0); + + __vector_quad acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7; + __builtin_mma_xxsetaccz(&acc0); + __builtin_mma_xxsetaccz(&acc1); + __builtin_mma_xxsetaccz(&acc2); + __builtin_mma_xxsetaccz(&acc3); + __builtin_mma_xxsetaccz(&acc4); + __builtin_mma_xxsetaccz(&acc5); + __builtin_mma_xxsetaccz(&acc6); + __builtin_mma_xxsetaccz(&acc7); + + const int8_t* __restrict__ ap = a_packed; + const uint8_t* __restrict__ bp = b_panel; + +#pragma GCC unroll 4 + for (int32_t k = 0; k < k_size; k += 4) { + const __vector unsigned char avec = + vec_xl(0, reinterpret_cast(ap)); + __builtin_mma_xvi8ger4pp(&acc0, avec, vec_xl(0, bp)); + __builtin_mma_xvi8ger4pp(&acc1, avec, vec_xl(16, bp)); + __builtin_mma_xvi8ger4pp(&acc2, avec, vec_xl(32, bp)); + __builtin_mma_xvi8ger4pp(&acc3, avec, vec_xl(48, bp)); + __builtin_mma_xvi8ger4pp(&acc4, avec, vec_xl(64, bp)); + __builtin_mma_xvi8ger4pp(&acc5, avec, vec_xl(80, bp)); + __builtin_mma_xvi8ger4pp(&acc6, avec, vec_xl(96, bp)); + __builtin_mma_xvi8ger4pp(&acc7, avec, vec_xl(112, bp)); + ap += 16; + bp += 128; + } + + __vector signed int res[4]; +#define STORE_ACC(ACC, CG_IDX) \ + { \ + __builtin_mma_disassemble_acc(res, &ACC); \ + for (int r = 0; r < m; ++r) { \ + const __vector signed int vcorr = vec_splats(128 * a_sums[r]); \ + const __vector signed int val = vec_sub(res[r], vcorr); \ + vec_xst(val, 0, \ + reinterpret_cast(c_ptr + r * ldc + (CG_IDX) * 4)); \ + } \ + } + STORE_ACC(acc0, 0); + STORE_ACC(acc1, 1); + STORE_ACC(acc2, 2); + STORE_ACC(acc3, 3); + STORE_ACC(acc4, 4); + STORE_ACC(acc5, 5); + STORE_ACC(acc6, 6); + STORE_ACC(acc7, 7); +#undef STORE_ACC +} + +} // namespace vsx_int8 + +// ============================================================================= +// MicroGemmINT8 specialisation for VSX (POWER10+) +// ============================================================================= +template +class MicroGemmINT8 { + public: + static constexpr int32_t K = vsx_int8::K; // 4 + static constexpr int32_t Mr = vsx_int8::Mr; // 4 + static constexpr int32_t Nr = vsx_int8::Nr; // 32 + static constexpr int32_t NrGemv = Nr; + static constexpr int32_t MaxMSize = 8; + static constexpr int32_t NSize = Nr; + static constexpr int32_t WeightOCGroupSize = Nr; + static_assert(MaxMSize % Mr == 0); + + // ------------------------------------------------------------------------- + // quantize_row: scalar_t -> INT8 + per-token scale + // Single-pass 16-element vector conversion with aligned 16-byte stores. + // ------------------------------------------------------------------------- + static FORCE_INLINE void quantize_row(const scalar_t* input, int8_t* output, + float& scale, const int32_t size) { + TORCH_CHECK_EQ(size % 16, 0); + + const __vector signed short z = vec_splats((signed short)0); + + auto load_lo_hi = + [&](int32_t i) -> std::pair<__vector float, __vector float> { + if constexpr (std::is_same_v) { + const float* p = reinterpret_cast(input) + i; + return {vec_xl(0, p), vec_xl(16, p)}; + } else { + const signed short* p = + reinterpret_cast(input + i); + const __vector signed short raw = (__vector signed short)vec_xl(0, p); + return {(__vector float)vec_mergeh(z, raw), + (__vector float)vec_mergel(z, raw)}; + } + }; + + __vector float max_v = vec_splats(0.0f); + for (int32_t i = 0; i < size; i += 8) { + auto [lo, hi] = load_lo_hi(i); + max_v = vec_max(max_v, vec_abs(lo)); + max_v = vec_max(max_v, vec_abs(hi)); + } + __vector float t = vec_max(max_v, vec_sld(max_v, max_v, 8)); + t = vec_max(t, vec_sld(t, t, 4)); + const float abs_max = std::max(vec_extract(t, 0), 1.0e-7f); + scale = abs_max / 127.0f; + const __vector float inv_scale = vec_splats(127.0f / abs_max); + + for (int32_t i = 0; i < size; i += 16) { + auto [lo0, hi0] = load_lo_hi(i); + auto [lo1, hi1] = load_lo_hi(i + 8); + + const __vector signed int i0 = + vec_cts(vec_round(vec_mul(lo0, inv_scale)), 0); + const __vector signed int i1 = + vec_cts(vec_round(vec_mul(hi0, inv_scale)), 0); + const __vector signed int i2 = + vec_cts(vec_round(vec_mul(lo1, inv_scale)), 0); + const __vector signed int i3 = + vec_cts(vec_round(vec_mul(hi1, inv_scale)), 0); + + const __vector signed short s0 = vec_packs(i0, i1); + const __vector signed short s1 = vec_packs(i2, i3); + const __vector signed char s8 = vec_packs(s0, s1); + + vec_xst(s8, 0, reinterpret_cast(output + i)); + } + } + + // ------------------------------------------------------------------------- + // dequantize_tile: INT32 accumulator -> FP32 scaled output + // ------------------------------------------------------------------------- + static FORCE_INLINE void dequantize_tile( + int32_t* input, float* output, const float* __restrict__ input_scales, + const float* __restrict__ weight_scales, const int32_t m, const int32_t n, + const int32_t stride) { + TORCH_CHECK_EQ(n % 4, 0); + for (int32_t mi = 0; mi < m; ++mi) { + const __vector float a_scale = vec_splats(input_scales[mi]); + for (int32_t ni = 0; ni < n; ni += 4) { + const __vector signed int i32 = vec_xl(0, input + mi * stride + ni); + const __vector float wscale = vec_xl(0, weight_scales + ni); + vec_xst(vec_mul(vec_ctf(i32, 0), vec_mul(a_scale, wscale)), 0, + output + mi * stride + ni); + } + } + } + + // ------------------------------------------------------------------------- + // pack_input_from_rows: pack rows into 4-row MMA tiles and compute row sums + // Header: 64 bytes (contains int32 a_sums[0..7]) + // Data: a_packed + 64 (contains 4x4 tiles of A) + // ------------------------------------------------------------------------- + static void pack_input_from_rows(const int8_t* const* __restrict__ rows, + int8_t* __restrict__ a_packed, + const int32_t m, const int32_t k) { + TORCH_CHECK(m > 0 && m <= MaxMSize); + TORCH_CHECK_EQ(k % K, 0); + + int32_t* a_sums = reinterpret_cast(a_packed); + for (int32_t r = 0; r < 8; ++r) a_sums[r] = 0; + + int8_t* a_data = a_packed + 64; + + for (int32_t row_base = 0; row_base < m; row_base += Mr) { + const int32_t panel_m = std::min(Mr, m - row_base); + const int8_t* const* panel_rows = rows + row_base; + int32_t* panel_sums = a_sums + row_base; + int8_t* dst_panel = a_data + (row_base / Mr) * (k * Mr); + + for (int32_t ki = 0; ki < k; ki += 4) { + int8_t* dst = dst_panel + (ki / 4) * 16; + for (int32_t r = 0; r < 4; ++r) { + if (r < panel_m) { + const int8_t* src = panel_rows[r] + ki; + dst[r * 4 + 0] = src[0]; + dst[r * 4 + 1] = src[1]; + dst[r * 4 + 2] = src[2]; + dst[r * 4 + 3] = src[3]; + panel_sums[r] += + static_cast(src[0]) + src[1] + src[2] + src[3]; + } else { + dst[r * 4 + 0] = 0; + dst[r * 4 + 1] = 0; + dst[r * 4 + 2] = 0; + dst[r * 4 + 3] = 0; + } + } + } + } + } + + // ------------------------------------------------------------------------- + // pack_weight: [output_size, input_size] INT8 -> packed layout + // Pre-biases with +128 in [N/32][K/4][8 col-groups][4 cols x 4 k] layout + // ------------------------------------------------------------------------- + static void pack_weight(const int8_t* __restrict__ weight, + int8_t* __restrict__ packed_weight, + const int32_t output_size, const int32_t input_size) { + TORCH_CHECK_EQ(output_size % NSize, 0); + TORCH_CHECK_EQ(input_size % K, 0); + + uint8_t* dst_base = reinterpret_cast(packed_weight); + + for (int32_t o = 0; o < output_size; o += Nr) { + uint8_t* dst_panel = dst_base + o * input_size; + for (int32_t k = 0; k < input_size; k += 4) { + for (int32_t cg = 0; cg < 8; ++cg) { + uint8_t* dst = dst_panel + (k / 4) * 128 + cg * 16; + for (int32_t col = 0; col < 4; ++col) { + for (int32_t ki = 0; ki < 4; ++ki) { + const int8_t w = weight[(o + cg * 4 + col) * input_size + k + ki]; + dst[col * 4 + ki] = + static_cast(static_cast(w) + 128); + } + } + } + } + } + } + + // ------------------------------------------------------------------------- + // gemm: 8-accumulator MMA GEMM over 32-column tiles + // ------------------------------------------------------------------------- + void gemm(const int8_t* __restrict__ a_packed, + const int8_t* __restrict__ b_packed, int32_t* __restrict__ c, + const int32_t m, const int32_t k, const int64_t b_n_group_stride, + const int64_t ldc) const { + TORCH_CHECK(m > 0 && m <= MaxMSize); + TORCH_CHECK_EQ(k % K, 0); + (void)b_n_group_stride; + + const int32_t* a_sums = reinterpret_cast(a_packed); + const int8_t* a_data = a_packed + 64; + const uint8_t* b_u8 = reinterpret_cast(b_packed); + + for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr) { + const uint8_t* b_panel = b_u8 + n_idx * k; + + for (int32_t row_base = 0; row_base < m; row_base += Mr) { + const int32_t panel_m = std::min(Mr, m - row_base); + const int8_t* a_panel = a_data + (row_base / Mr) * (k * Mr); + const int32_t* panel_sums = a_sums + row_base; + + vsx_int8::gemm_mma_4x32(a_panel, b_panel, panel_sums, + c + row_base * ldc + n_idx, panel_m, k, ldc); + } + } + } +}; + +} // namespace cpu_micro_gemm + +#endif // CPU_MICRO_GEMM_INT8_VSX_HPP diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 8f663a93c3..88170068de 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -883,8 +883,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "bool skip_weighted, " "str act, str isa) -> ()"); ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe); -#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \ - !defined(__APPLE__) + +#if (defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \ + !defined(__APPLE__)) || \ + defined(__powerpc64__) ops.def( "prepack_moe_weight_int8(Tensor weight, Tensor(a1!) packed_weight, " "str isa) -> ()"); @@ -895,8 +897,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, bool " "skip_weighted, str act, str isa) -> ()"); ops.impl("cpu_fused_moe_int8", torch::kCPU, &cpu_fused_moe_int8); -#endif // #if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && - // !defined(__APPLE__) +#endif // ARM_I8MM+BF16 or __powerpc64__ ops.def( "mla_decode_kvcache(" " Tensor! out, Tensor query, Tensor kv_cache," diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py index cf83ef6fc0..d4c5437221 100644 --- a/tests/kernels/moe/test_cpu_fused_moe.py +++ b/tests/kernels/moe/test_cpu_fused_moe.py @@ -306,17 +306,28 @@ def test_cpu_fused_moe( @pytest.mark.skipif( - current_platform.get_cpu_architecture() != CpuArchEnum.ARM, - reason="Requires Arm CPU", + current_platform.get_cpu_architecture() + not in (CpuArchEnum.ARM, CpuArchEnum.POWERPC), + reason="Requires Arm or POWER CPU", ) @pytest.mark.parametrize("batch_size", BATCH_SIZE) @pytest.mark.parametrize("expert_num", EXPERT_NUM) @pytest.mark.parametrize("hidden_size", HIDDEN_DIM) @pytest.mark.parametrize("intermediate_size", INTERMEDIATE_DIM) @pytest.mark.parametrize("use_bias", USE_BIAS) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "dtype", + ( + [torch.float32, torch.float16, torch.bfloat16] + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM + else [torch.float32, torch.bfloat16] + ), +) @pytest.mark.parametrize("act", ACT) -@pytest.mark.parametrize("isa", ["neon"]) +@pytest.mark.parametrize( + "isa", + ["neon"] if current_platform.get_cpu_architecture() == CpuArchEnum.ARM else ["vsx"], +) def test_cpu_fused_moe_int8( batch_size: int, expert_num: int, @@ -386,7 +397,8 @@ def test_cpu_fused_moe_int8( isa, ) - torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=2e-2) + atol = rtol = 1e-1 if isa == "vsx" else 2e-2 + torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol) # moe_intermediate_size not a multiple of 32, e.g. what tensor-parallel 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 7b3193fb5f..858b351864 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -950,6 +950,29 @@ class CPUExpertsInt4(mk.FusedMoEExpertsModular): # =========================================================================== +def prepare_int8_moe_layer_for_cpu( + w13: torch.Tensor, + w2: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Prepack INT8 MoE weights for the current CPU architecture.""" + # SMMLA packing for AArch64 + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + return ( + cpu_prepack_moe_weight_int8(w13, "neon"), + cpu_prepack_moe_weight_int8(w2, "neon"), + ) + # VSX packing for POWER + if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC: + return ( + cpu_prepack_moe_weight_int8(w13, "vsx"), + cpu_prepack_moe_weight_int8(w2, "vsx"), + ) + # VNNI packing for x86 + packed_w13 = torch.ops._C.convert_weight_packed(w13) + packed_w2 = torch.ops._C.convert_weight_packed(w2) + return packed_w13, packed_w2 + + class CPUExpertsInt8(mk.FusedMoEExpertsModular): """CPU INT8 W8A8 per-channel weight / dynamic per-token activation modular MoE experts.""" @@ -1283,6 +1306,160 @@ class ArmCPUExpertsInt8(mk.FusedMoEExpertsModular): ) +class PowerCPUExpertsInt8(mk.FusedMoEExpertsModular): + """POWER VSX INT8 MoE with per-token activation and channelwise weight quant.""" + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @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.FusedMoEExpertsModular.is_supported_config( + cls, + moe_config, + weight_key, + activation_key, + activation_format, + ) + if not supported: + return supported, reason + if moe_config.in_dtype not in (torch.float32, torch.bfloat16): + return False, "kernel requires float32 or bfloat16 activations" + if moe_config.hidden_dim % 16 != 0: + return False, "kernel requires hidden dim divisible by 16" + if moe_config.intermediate_size_per_partition % 16 != 0: + return False, "kernel requires intermediate dim divisible by 16" + return True, None + + @staticmethod + def _supports_current_device() -> bool: + return ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC + and hasattr(torch.ops._C, "cpu_fused_moe_int8") + ) + + @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_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return not moe_parallel_config.use_ep + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == ( + kInt8StaticChannelSym, + kInt8DynamicTokenSym, + ) + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Default, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + + @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: + w13 = cpu_prepack_moe_weight_int8(layer.w13_weight, "vsx") + w2 = cpu_prepack_moe_weight_int8(layer.w2_weight, "vsx") + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # cpu_fused_moe_int8 manages its own scratch space. + return (0,), (0,), (M, K) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ) -> None: + # The modular prepare step applies router weights to the input when + # apply_router_weight_on_input is enabled. + assert self.w1_scale is not None + assert self.w2_scale is not None + cpu_fused_moe_int8( + output, + hidden_states, + w1, + w2, + self.w1_scale, + self.w2_scale, + self.w1_bias, + self.w2_bias, + topk_weights, + topk_ids, + activation.value, + "vsx", + skip_weighted=apply_router_weight_on_input, + ) + + class ZenCPUExpertsInt8(mk.FusedMoEExpertsModular): """AMD Zen INT8 MoE with per-token activation and channelwise weight quantization, dispatched through zentorch.""" diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index 3b8a3c4f87..1a23b01736 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -71,10 +71,16 @@ def backend_to_kernel_cls( from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( ArmCPUExpertsInt8, CPUExpertsInt8, + PowerCPUExpertsInt8, ZenCPUExpertsInt8, ) - return [ZenCPUExpertsInt8, ArmCPUExpertsInt8, CPUExpertsInt8] + return [ + ZenCPUExpertsInt8, + PowerCPUExpertsInt8, + ArmCPUExpertsInt8, + CPUExpertsInt8, + ] else: raise ValueError(f"Unknown Int8 MoE backend: {backend.value}")