iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[HARDWARE][POWER] Enable W8A8 INT8 MoE on POWER (#55316)

Browse Source 
Signed-off-by: Rukhaiya <bibirukhaiya123@gmail.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Co-authored-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Co-authored-by: Antigravity <antigravity@google.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
This commit is contained in:
Rukhaiya2004
2026-09-17 21:45:13 +08:00
committed by GitHub
co-authored by Akash kaothalkar Antigravity Li, Jiang <jiang1.li@intel.com>
parent 0eae9acd4d
commit 117cf43fb6
8 changed files with 537 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmake/cpu_extension.cmake
+6
View File
@@ -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"
csrc/cpu/cpu_fused_moe_int8.cpp
+16 -3
View File
@@ -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<cpu_utils::ISA::VSX, SCALAR_TYPE>; \
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();
csrc/cpu/cpu_types_vsx.hpp
+22 -9
View File
@@ -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<FP32Vec8> {
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<FP32Vec16> {
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<FP32Vec16> {
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<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]),
csrc/cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp
+287
View File
@@ -0,0 +1,287 @@
#ifndef CPU_MICRO_GEMM_INT8_VSX_HPP
#define CPU_MICRO_GEMM_INT8_VSX_HPP
#include <algorithm>
#include <cstdint>
#include <cstring>
#include <type_traits>
#include <utility>
#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
#include <altivec.h>
#include <c10/util/BFloat16.h>
#include <c10/util/Exception.h>
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<const unsigned char*>(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<signed int*>(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 <typename scalar_t>
class MicroGemmINT8<cpu_utils::ISA::VSX, scalar_t> {
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<scalar_t, float>) {
const float* p = reinterpret_cast<const float*>(input) + i;
return {vec_xl(0, p), vec_xl(16, p)};
} else {
const signed short* p =
reinterpret_cast<const signed short*>(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<signed char*>(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<int32_t*>(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<int32_t>(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<uint8_t*>(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<uint8_t>(static_cast<int>(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<const int32_t*>(a_packed);
const int8_t* a_data = a_packed + 64;
const uint8_t* b_u8 = reinterpret_cast<const uint8_t*>(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
csrc/cpu/torch_bindings.cpp
+5 -4
View File
@@ -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,"
tests/kernels/moe/test_cpu_fused_moe.py
+17 -5
View File
@@ -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
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+177
View File
@@ -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."""
vllm/model_executor/layers/fused_moe/oracle/int8.py
+7 -1
View File
@@ -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}")