[HARDWARE][POWER] Enable W8A8 INT8 MoE on POWER (#55316)
[HARDWARE][POWER] Enable W8A8 INT8 MoE on POWER (#55316) 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>
Rukhaiya2004 committed
Sep 17, 2026 at 19:15 UTC
117cf43fb65713a0a7a3c35b23be078d5fb319d4
8 files changed
+537
-22
cmake/cpu_extension.cmake
+6
@@ -522,6 +522,12 @@ if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
522
${VLLM_EXT_SRC})
523
endif()
524
525
+if (POWER10_FOUND OR POWER11_FOUND)
526
+ set(VLLM_EXT_SRC
527
+ "csrc/cpu/cpu_fused_moe_int8.cpp"
528
+ ${VLLM_EXT_SRC})
529
+endif()
530
+
531
if(USE_ONEDNN)
532
set(VLLM_EXT_SRC
533
"csrc/cpu/dnnl_kernels.cpp"
csrc/cpu/cpu_fused_moe_int8.cpp
+16
-3
@@ -26,10 +26,23 @@
26
#define NEON_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::NEON:
27
#endif
28
29
+#if defined(__powerpc64__)
30
+ #include "cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp"
31
+ #define VSX_DISPATCH(SCALAR_TYPE, ...) \
32
+ case cpu_utils::ISA::VSX: { \
33
+ using gemm_t = \
34
+ cpu_micro_gemm::MicroGemmINT8<cpu_utils::ISA::VSX, SCALAR_TYPE>; \
35
+ return __VA_ARGS__(); \
36
+ }
37
+#else
38
+ #define VSX_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::VSX:
39
+#endif
40
+
41
#define CPU_INT8_ISA_DISPATCH_IMPL(ISA_TYPE, SCALAR_TYPE, ...) \
42
[&] { \
43
switch (ISA_TYPE) { \
44
NEON_DISPATCH(SCALAR_TYPE, __VA_ARGS__) \
45
+ VSX_DISPATCH(SCALAR_TYPE, __VA_ARGS__) \
46
default: { \
47
TORCH_CHECK(false, "Invalid CPU ISA type."); \
48
} \
@@ -81,9 +94,9 @@ void fused_moe_int8_impl(
94
95
const int32_t thread_num = cpu_utils::get_max_threads();
96
const int32_t w13_input_buffer_size = cpu_utils::round_up<64>(
84
- gemm_m_tile_size * input_size_13 * sizeof(int8_t));
85
- const int32_t w2_input_buffer_size =
86
- cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 * sizeof(int8_t));
97
+ (gemm_m_tile_size * input_size_13 + 64) * sizeof(int8_t));
98
+ const int32_t w2_input_buffer_size = cpu_utils::round_up<64>(
99
+ (gemm_m_tile_size * input_size_2 + 64) * sizeof(int8_t));
100
101
const int32_t w13_n_tile_size = [&]() {
102
const int64_t cache_size = cpu_utils::get_available_l2_size();
csrc/cpu/cpu_types_vsx.hpp
+22
-9
@@ -13,10 +13,10 @@ namespace vec_op {
13
struct fp8_e4m3_tag {};
14
struct fp8_e5m2_tag {};
15
16
-#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
17
- AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
18
- AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
19
- AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
16
+#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
17
+ AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
18
+ AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
19
+// Note: FP16 (Half) is not supported on POWER VSX — no FP16Vec16 type.
20
21
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
22
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
@@ -416,6 +416,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
416
tmp.val[0] = two_x;
417
tmp.val[1] = two_x;
418
FP32Vec8 temp_vec(tmp);
419
+ FP32Vec8 exp_vec = temp_vec.exp();
420
vector float e = temp_vec.exp().reg.val[0];
421
422
vector float num = vec_sub(e, one);
@@ -634,6 +635,16 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
635
reg.val[3] = vec_ctf(v.reg.val[3], 0);
636
}
637
638
+ explicit FP32Vec16(const float* __restrict__ base, const INT32Vec16& index) {
639
+ INT32Vec16::AliasReg idx;
640
+ idx.reg = index.reg;
641
+ AliasReg ar;
642
+ for (int i = 0; i < VEC_ELEM_NUM; ++i) {
643
+ ar.values[i] = base[idx.values[i]];
644
+ }
645
+ reg = ar.reg;
646
+ }
647
+
648
FP32Vec16 operator*(const FP32Vec16& b) const {
649
return FP32Vec16(f32x4x4_t({vec_mul(reg.val[0], b.reg.val[0]),
650
vec_mul(reg.val[1], b.reg.val[1]),
@@ -648,6 +659,13 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
659
vec_add(reg.val[3], b.reg.val[3])}));
660
}
661
662
+ FP32Vec16 operator-() const {
663
+ const __vector float zero = vec_splats(0.0f);
664
+ return FP32Vec16(
665
+ f32x4x4_t({vec_sub(zero, reg.val[0]), vec_sub(zero, reg.val[1]),
666
+ vec_sub(zero, reg.val[2]), vec_sub(zero, reg.val[3])}));
667
+ }
668
+
669
FP32Vec16 operator-(const FP32Vec16& b) const {
670
return FP32Vec16(f32x4x4_t({vec_sub(reg.val[0], b.reg.val[0]),
671
vec_sub(reg.val[1], b.reg.val[1]),
@@ -655,11 +673,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
673
vec_sub(reg.val[3], b.reg.val[3])}));
674
}
675
658
- FP32Vec16 operator-() const {
659
- return FP32Vec16(
660
- f32x4x4_t({-reg.val[0], -reg.val[1], -reg.val[2], -reg.val[3]}));
661
- }
662
-
676
FP32Vec16 operator/(const FP32Vec16& b) const {
677
return FP32Vec16(f32x4x4_t({vec_div(reg.val[0], b.reg.val[0]),
678
vec_div(reg.val[1], b.reg.val[1]),
csrc/cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp
new
+287
@@ -0,0 +1,287 @@
1
+#ifndef CPU_MICRO_GEMM_INT8_VSX_HPP
2
+#define CPU_MICRO_GEMM_INT8_VSX_HPP
3
+
4
+#include <algorithm>
5
+#include <cstdint>
6
+#include <cstring>
7
+#include <type_traits>
8
+#include <utility>
9
+
10
+#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
11
+
12
+#include <altivec.h>
13
+#include <c10/util/BFloat16.h>
14
+#include <c10/util/Exception.h>
15
+
16
+namespace cpu_micro_gemm {
17
+
18
+namespace vsx_int8 {
19
+
20
+constexpr int32_t K = 4;
21
+constexpr int32_t Mr = 4;
22
+constexpr int32_t Nr = 32;
23
+
24
+// ---------------------------------------------------------------------------
25
+// 8-Accumulator MMA 4x32 micro-GEMM (POWER10 xvi8ger4pp).
26
+//
27
+// Computes 4 rows x 32 columns (128 outputs) using all 8 hardware accumulators.
28
+// Operates on pre-biased uint8 weights (B + 128) and subtracts 128 * a_sums.
29
+// Works for all m in [1, 4] with zero scalar loop overhead.
30
+// ---------------------------------------------------------------------------
31
+__attribute__((target("cpu=power10"))) FORCE_INLINE void gemm_mma_4x32(
32
+ const int8_t* __restrict__ a_packed, const uint8_t* __restrict__ b_panel,
33
+ const int32_t* __restrict__ a_sums, int32_t* __restrict__ c_ptr,
34
+ const int32_t m, const int32_t k_size, const int64_t ldc) {
35
+ TORCH_CHECK(m > 0 && m <= Mr);
36
+ TORCH_CHECK(k_size % K == 0);
37
+
38
+ __vector_quad acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7;
39
+ __builtin_mma_xxsetaccz(&acc0);
40
+ __builtin_mma_xxsetaccz(&acc1);
41
+ __builtin_mma_xxsetaccz(&acc2);
42
+ __builtin_mma_xxsetaccz(&acc3);
43
+ __builtin_mma_xxsetaccz(&acc4);
44
+ __builtin_mma_xxsetaccz(&acc5);
45
+ __builtin_mma_xxsetaccz(&acc6);
46
+ __builtin_mma_xxsetaccz(&acc7);
47
+
48
+ const int8_t* __restrict__ ap = a_packed;
49
+ const uint8_t* __restrict__ bp = b_panel;
50
+
51
+#pragma GCC unroll 4
52
+ for (int32_t k = 0; k < k_size; k += 4) {
53
+ const __vector unsigned char avec =
54
+ vec_xl(0, reinterpret_cast<const unsigned char*>(ap));
55
+ __builtin_mma_xvi8ger4pp(&acc0, avec, vec_xl(0, bp));
56
+ __builtin_mma_xvi8ger4pp(&acc1, avec, vec_xl(16, bp));
57
+ __builtin_mma_xvi8ger4pp(&acc2, avec, vec_xl(32, bp));
58
+ __builtin_mma_xvi8ger4pp(&acc3, avec, vec_xl(48, bp));
59
+ __builtin_mma_xvi8ger4pp(&acc4, avec, vec_xl(64, bp));
60
+ __builtin_mma_xvi8ger4pp(&acc5, avec, vec_xl(80, bp));
61
+ __builtin_mma_xvi8ger4pp(&acc6, avec, vec_xl(96, bp));
62
+ __builtin_mma_xvi8ger4pp(&acc7, avec, vec_xl(112, bp));
63
+ ap += 16;
64
+ bp += 128;
65
+ }
66
+
67
+ __vector signed int res[4];
68
+#define STORE_ACC(ACC, CG_IDX) \
69
+ { \
70
+ __builtin_mma_disassemble_acc(res, &ACC); \
71
+ for (int r = 0; r < m; ++r) { \
72
+ const __vector signed int vcorr = vec_splats(128 * a_sums[r]); \
73
+ const __vector signed int val = vec_sub(res[r], vcorr); \
74
+ vec_xst(val, 0, \
75
+ reinterpret_cast<signed int*>(c_ptr + r * ldc + (CG_IDX) * 4)); \
76
+ } \
77
+ }
78
+ STORE_ACC(acc0, 0);
79
+ STORE_ACC(acc1, 1);
80
+ STORE_ACC(acc2, 2);
81
+ STORE_ACC(acc3, 3);
82
+ STORE_ACC(acc4, 4);
83
+ STORE_ACC(acc5, 5);
84
+ STORE_ACC(acc6, 6);
85
+ STORE_ACC(acc7, 7);
86
+#undef STORE_ACC
87
+}
88
+
89
+} // namespace vsx_int8
90
+
91
+// =============================================================================
92
+// MicroGemmINT8 specialisation for VSX (POWER10+)
93
+// =============================================================================
94
+template <typename scalar_t>
95
+class MicroGemmINT8<cpu_utils::ISA::VSX, scalar_t> {
96
+ public:
97
+ static constexpr int32_t K = vsx_int8::K; // 4
98
+ static constexpr int32_t Mr = vsx_int8::Mr; // 4
99
+ static constexpr int32_t Nr = vsx_int8::Nr; // 32
100
+ static constexpr int32_t NrGemv = Nr;
101
+ static constexpr int32_t MaxMSize = 8;
102
+ static constexpr int32_t NSize = Nr;
103
+ static constexpr int32_t WeightOCGroupSize = Nr;
104
+ static_assert(MaxMSize % Mr == 0);
105
+
106
+ // -------------------------------------------------------------------------
107
+ // quantize_row: scalar_t -> INT8 + per-token scale
108
+ // Single-pass 16-element vector conversion with aligned 16-byte stores.
109
+ // -------------------------------------------------------------------------
110
+ static FORCE_INLINE void quantize_row(const scalar_t* input, int8_t* output,
111
+ float& scale, const int32_t size) {
112
+ TORCH_CHECK_EQ(size % 16, 0);
113
+
114
+ const __vector signed short z = vec_splats((signed short)0);
115
+
116
+ auto load_lo_hi =
117
+ [&](int32_t i) -> std::pair<__vector float, __vector float> {
118
+ if constexpr (std::is_same_v<scalar_t, float>) {
119
+ const float* p = reinterpret_cast<const float*>(input) + i;
120
+ return {vec_xl(0, p), vec_xl(16, p)};
121
+ } else {
122
+ const signed short* p =
123
+ reinterpret_cast<const signed short*>(input + i);
124
+ const __vector signed short raw = (__vector signed short)vec_xl(0, p);
125
+ return {(__vector float)vec_mergeh(z, raw),
126
+ (__vector float)vec_mergel(z, raw)};
127
+ }
128
+ };
129
+
130
+ __vector float max_v = vec_splats(0.0f);
131
+ for (int32_t i = 0; i < size; i += 8) {
132
+ auto [lo, hi] = load_lo_hi(i);
133
+ max_v = vec_max(max_v, vec_abs(lo));
134
+ max_v = vec_max(max_v, vec_abs(hi));
135
+ }
136
+ __vector float t = vec_max(max_v, vec_sld(max_v, max_v, 8));
137
+ t = vec_max(t, vec_sld(t, t, 4));
138
+ const float abs_max = std::max(vec_extract(t, 0), 1.0e-7f);
139
+ scale = abs_max / 127.0f;
140
+ const __vector float inv_scale = vec_splats(127.0f / abs_max);
141
+
142
+ for (int32_t i = 0; i < size; i += 16) {
143
+ auto [lo0, hi0] = load_lo_hi(i);
144
+ auto [lo1, hi1] = load_lo_hi(i + 8);
145
+
146
+ const __vector signed int i0 =
147
+ vec_cts(vec_round(vec_mul(lo0, inv_scale)), 0);
148
+ const __vector signed int i1 =
149
+ vec_cts(vec_round(vec_mul(hi0, inv_scale)), 0);
150
+ const __vector signed int i2 =
151
+ vec_cts(vec_round(vec_mul(lo1, inv_scale)), 0);
152
+ const __vector signed int i3 =
153
+ vec_cts(vec_round(vec_mul(hi1, inv_scale)), 0);
154
+
155
+ const __vector signed short s0 = vec_packs(i0, i1);
156
+ const __vector signed short s1 = vec_packs(i2, i3);
157
+ const __vector signed char s8 = vec_packs(s0, s1);
158
+
159
+ vec_xst(s8, 0, reinterpret_cast<signed char*>(output + i));
160
+ }
161
+ }
162
+
163
+ // -------------------------------------------------------------------------
164
+ // dequantize_tile: INT32 accumulator -> FP32 scaled output
165
+ // -------------------------------------------------------------------------
166
+ static FORCE_INLINE void dequantize_tile(
167
+ int32_t* input, float* output, const float* __restrict__ input_scales,
168
+ const float* __restrict__ weight_scales, const int32_t m, const int32_t n,
169
+ const int32_t stride) {
170
+ TORCH_CHECK_EQ(n % 4, 0);
171
+ for (int32_t mi = 0; mi < m; ++mi) {
172
+ const __vector float a_scale = vec_splats(input_scales[mi]);
173
+ for (int32_t ni = 0; ni < n; ni += 4) {
174
+ const __vector signed int i32 = vec_xl(0, input + mi * stride + ni);
175
+ const __vector float wscale = vec_xl(0, weight_scales + ni);
176
+ vec_xst(vec_mul(vec_ctf(i32, 0), vec_mul(a_scale, wscale)), 0,
177
+ output + mi * stride + ni);
178
+ }
179
+ }
180
+ }
181
+
182
+ // -------------------------------------------------------------------------
183
+ // pack_input_from_rows: pack rows into 4-row MMA tiles and compute row sums
184
+ // Header: 64 bytes (contains int32 a_sums[0..7])
185
+ // Data: a_packed + 64 (contains 4x4 tiles of A)
186
+ // -------------------------------------------------------------------------
187
+ static void pack_input_from_rows(const int8_t* const* __restrict__ rows,
188
+ int8_t* __restrict__ a_packed,
189
+ const int32_t m, const int32_t k) {
190
+ TORCH_CHECK(m > 0 && m <= MaxMSize);
191
+ TORCH_CHECK_EQ(k % K, 0);
192
+
193
+ int32_t* a_sums = reinterpret_cast<int32_t*>(a_packed);
194
+ for (int32_t r = 0; r < 8; ++r) a_sums[r] = 0;
195
+
196
+ int8_t* a_data = a_packed + 64;
197
+
198
+ for (int32_t row_base = 0; row_base < m; row_base += Mr) {
199
+ const int32_t panel_m = std::min(Mr, m - row_base);
200
+ const int8_t* const* panel_rows = rows + row_base;
201
+ int32_t* panel_sums = a_sums + row_base;
202
+ int8_t* dst_panel = a_data + (row_base / Mr) * (k * Mr);
203
+
204
+ for (int32_t ki = 0; ki < k; ki += 4) {
205
+ int8_t* dst = dst_panel + (ki / 4) * 16;
206
+ for (int32_t r = 0; r < 4; ++r) {
207
+ if (r < panel_m) {
208
+ const int8_t* src = panel_rows[r] + ki;
209
+ dst[r * 4 + 0] = src[0];
210
+ dst[r * 4 + 1] = src[1];
211
+ dst[r * 4 + 2] = src[2];
212
+ dst[r * 4 + 3] = src[3];
213
+ panel_sums[r] +=
214
+ static_cast<int32_t>(src[0]) + src[1] + src[2] + src[3];
215
+ } else {
216
+ dst[r * 4 + 0] = 0;
217
+ dst[r * 4 + 1] = 0;
218
+ dst[r * 4 + 2] = 0;
219
+ dst[r * 4 + 3] = 0;
220
+ }
221
+ }
222
+ }
223
+ }
224
+ }
225
+
226
+ // -------------------------------------------------------------------------
227
+ // pack_weight: [output_size, input_size] INT8 -> packed layout
228
+ // Pre-biases with +128 in [N/32][K/4][8 col-groups][4 cols x 4 k] layout
229
+ // -------------------------------------------------------------------------
230
+ static void pack_weight(const int8_t* __restrict__ weight,
231
+ int8_t* __restrict__ packed_weight,
232
+ const int32_t output_size, const int32_t input_size) {
233
+ TORCH_CHECK_EQ(output_size % NSize, 0);
234
+ TORCH_CHECK_EQ(input_size % K, 0);
235
+
236
+ uint8_t* dst_base = reinterpret_cast<uint8_t*>(packed_weight);
237
+
238
+ for (int32_t o = 0; o < output_size; o += Nr) {
239
+ uint8_t* dst_panel = dst_base + o * input_size;
240
+ for (int32_t k = 0; k < input_size; k += 4) {
241
+ for (int32_t cg = 0; cg < 8; ++cg) {
242
+ uint8_t* dst = dst_panel + (k / 4) * 128 + cg * 16;
243
+ for (int32_t col = 0; col < 4; ++col) {
244
+ for (int32_t ki = 0; ki < 4; ++ki) {
245
+ const int8_t w = weight[(o + cg * 4 + col) * input_size + k + ki];
246
+ dst[col * 4 + ki] =
247
+ static_cast<uint8_t>(static_cast<int>(w) + 128);
248
+ }
249
+ }
250
+ }
251
+ }
252
+ }
253
+ }
254
+
255
+ // -------------------------------------------------------------------------
256
+ // gemm: 8-accumulator MMA GEMM over 32-column tiles
257
+ // -------------------------------------------------------------------------
258
+ void gemm(const int8_t* __restrict__ a_packed,
259
+ const int8_t* __restrict__ b_packed, int32_t* __restrict__ c,
260
+ const int32_t m, const int32_t k, const int64_t b_n_group_stride,
261
+ const int64_t ldc) const {
262
+ TORCH_CHECK(m > 0 && m <= MaxMSize);
263
+ TORCH_CHECK_EQ(k % K, 0);
264
+ (void)b_n_group_stride;
265
+
266
+ const int32_t* a_sums = reinterpret_cast<const int32_t*>(a_packed);
267
+ const int8_t* a_data = a_packed + 64;
268
+ const uint8_t* b_u8 = reinterpret_cast<const uint8_t*>(b_packed);
269
+
270
+ for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr) {
271
+ const uint8_t* b_panel = b_u8 + n_idx * k;
272
+
273
+ for (int32_t row_base = 0; row_base < m; row_base += Mr) {
274
+ const int32_t panel_m = std::min(Mr, m - row_base);
275
+ const int8_t* a_panel = a_data + (row_base / Mr) * (k * Mr);
276
+ const int32_t* panel_sums = a_sums + row_base;
277
+
278
+ vsx_int8::gemm_mma_4x32(a_panel, b_panel, panel_sums,
279
+ c + row_base * ldc + n_idx, panel_m, k, ldc);
280
+ }
281
+ }
282
+ }
283
+};
284
+
285
+} // namespace cpu_micro_gemm
286
+
287
+#endif // CPU_MICRO_GEMM_INT8_VSX_HPP
csrc/cpu/torch_bindings.cpp
+5
-4
@@ -883,8 +883,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
883
"bool skip_weighted, "
884
"str act, str isa) -> ()");
885
ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
886
-#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \
887
- !defined(__APPLE__)
886
+
887
+#if (defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \
888
+ !defined(__APPLE__)) || \
889
+ defined(__powerpc64__)
890
ops.def(
891
"prepack_moe_weight_int8(Tensor weight, Tensor(a1!) packed_weight, "
892
"str isa) -> ()");
@@ -895,8 +897,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
897
"Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, bool "
898
"skip_weighted, str act, str isa) -> ()");
899
ops.impl("cpu_fused_moe_int8", torch::kCPU, &cpu_fused_moe_int8);
898
-#endif // #if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) &&
899
- // !defined(__APPLE__)
900
+#endif // ARM_I8MM+BF16 or __powerpc64__
901
ops.def(
902
"mla_decode_kvcache("
903
" Tensor! out, Tensor query, Tensor kv_cache,"
tests/kernels/moe/test_cpu_fused_moe.py
+17
-5
@@ -306,17 +306,28 @@ def test_cpu_fused_moe(
306
307
308
@pytest.mark.skipif(
309
- current_platform.get_cpu_architecture() != CpuArchEnum.ARM,
310
- reason="Requires Arm CPU",
309
+ current_platform.get_cpu_architecture()
310
+ not in (CpuArchEnum.ARM, CpuArchEnum.POWERPC),
311
+ reason="Requires Arm or POWER CPU",
312
)
313
@pytest.mark.parametrize("batch_size", BATCH_SIZE)
314
@pytest.mark.parametrize("expert_num", EXPERT_NUM)
315
@pytest.mark.parametrize("hidden_size", HIDDEN_DIM)
316
@pytest.mark.parametrize("intermediate_size", INTERMEDIATE_DIM)
317
@pytest.mark.parametrize("use_bias", USE_BIAS)
317
-@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
318
+@pytest.mark.parametrize(
319
+ "dtype",
320
+ (
321
+ [torch.float32, torch.float16, torch.bfloat16]
322
+ if current_platform.get_cpu_architecture() == CpuArchEnum.ARM
323
+ else [torch.float32, torch.bfloat16]
324
+ ),
325
+)
326
@pytest.mark.parametrize("act", ACT)
319
-@pytest.mark.parametrize("isa", ["neon"])
327
+@pytest.mark.parametrize(
328
+ "isa",
329
+ ["neon"] if current_platform.get_cpu_architecture() == CpuArchEnum.ARM else ["vsx"],
330
+)
331
def test_cpu_fused_moe_int8(
332
batch_size: int,
333
expert_num: int,
@@ -386,7 +397,8 @@ def test_cpu_fused_moe_int8(
397
isa,
398
)
399
389
- torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=2e-2)
400
+ atol = rtol = 1e-1 if isa == "vsx" else 2e-2
401
+ torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
402
403
404
# moe_intermediate_size not a multiple of 32, e.g. what tensor-parallel
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+177
@@ -950,6 +950,29 @@ class CPUExpertsInt4(mk.FusedMoEExpertsModular):
950
# ===========================================================================
951
952
953
+def prepare_int8_moe_layer_for_cpu(
954
+ w13: torch.Tensor,
955
+ w2: torch.Tensor,
956
+) -> tuple[torch.Tensor, torch.Tensor]:
957
+ """Prepack INT8 MoE weights for the current CPU architecture."""
958
+ # SMMLA packing for AArch64
959
+ if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
960
+ return (
961
+ cpu_prepack_moe_weight_int8(w13, "neon"),
962
+ cpu_prepack_moe_weight_int8(w2, "neon"),
963
+ )
964
+ # VSX packing for POWER
965
+ if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC:
966
+ return (
967
+ cpu_prepack_moe_weight_int8(w13, "vsx"),
968
+ cpu_prepack_moe_weight_int8(w2, "vsx"),
969
+ )
970
+ # VNNI packing for x86
971
+ packed_w13 = torch.ops._C.convert_weight_packed(w13)
972
+ packed_w2 = torch.ops._C.convert_weight_packed(w2)
973
+ return packed_w13, packed_w2
974
+
975
+
976
class CPUExpertsInt8(mk.FusedMoEExpertsModular):
977
"""CPU INT8 W8A8 per-channel weight / dynamic per-token activation
978
modular MoE experts."""
@@ -1283,6 +1306,160 @@ class ArmCPUExpertsInt8(mk.FusedMoEExpertsModular):
1306
)
1307
1308
1309
+class PowerCPUExpertsInt8(mk.FusedMoEExpertsModular):
1310
+ """POWER VSX INT8 MoE with per-token activation and channelwise weight quant."""
1311
+
1312
+ @property
1313
+ def expects_unquantized_inputs(self) -> bool:
1314
+ return True
1315
+
1316
+ @staticmethod
1317
+ def activation_format() -> mk.FusedMoEActivationFormat:
1318
+ return mk.FusedMoEActivationFormat.Standard
1319
+
1320
+ @staticmethod
1321
+ def is_supported_config(
1322
+ cls: type[mk.FusedMoEExperts],
1323
+ moe_config: FusedMoEConfig,
1324
+ weight_key: QuantKey | None,
1325
+ activation_key: QuantKey | None,
1326
+ activation_format: mk.FusedMoEActivationFormat,
1327
+ ) -> tuple[bool, str | None]:
1328
+ supported, reason = mk.FusedMoEExpertsModular.is_supported_config(
1329
+ cls,
1330
+ moe_config,
1331
+ weight_key,
1332
+ activation_key,
1333
+ activation_format,
1334
+ )
1335
+ if not supported:
1336
+ return supported, reason
1337
+ if moe_config.in_dtype not in (torch.float32, torch.bfloat16):
1338
+ return False, "kernel requires float32 or bfloat16 activations"
1339
+ if moe_config.hidden_dim % 16 != 0:
1340
+ return False, "kernel requires hidden dim divisible by 16"
1341
+ if moe_config.intermediate_size_per_partition % 16 != 0:
1342
+ return False, "kernel requires intermediate dim divisible by 16"
1343
+ return True, None
1344
+
1345
+ @staticmethod
1346
+ def _supports_current_device() -> bool:
1347
+ return (
1348
+ current_platform.is_cpu()
1349
+ and current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC
1350
+ and hasattr(torch.ops._C, "cpu_fused_moe_int8")
1351
+ )
1352
+
1353
+ @staticmethod
1354
+ def _supports_no_act_and_mul() -> bool:
1355
+ return False
1356
+
1357
+ @staticmethod
1358
+ def _supports_activation(activation: MoEActivation) -> bool:
1359
+ return activation in (
1360
+ MoEActivation.SILU,
1361
+ MoEActivation.SWIGLUOAI,
1362
+ MoEActivation.GELU,
1363
+ MoEActivation.GELU_TANH,
1364
+ )
1365
+
1366
+ @staticmethod
1367
+ def _supports_parallel_config(
1368
+ moe_parallel_config: FusedMoEParallelConfig,
1369
+ ) -> bool:
1370
+ return not moe_parallel_config.use_ep
1371
+
1372
+ @staticmethod
1373
+ def _supports_quant_scheme(
1374
+ weight_key: QuantKey | None,
1375
+ activation_key: QuantKey | None,
1376
+ ) -> bool:
1377
+ return (weight_key, activation_key) == (
1378
+ kInt8StaticChannelSym,
1379
+ kInt8DynamicTokenSym,
1380
+ )
1381
+
1382
+ @staticmethod
1383
+ def _supports_routing_method(
1384
+ routing_method: RoutingMethodType,
1385
+ weight_key: QuantKey | None,
1386
+ activation_key: QuantKey | None,
1387
+ ) -> bool:
1388
+ return routing_method in [
1389
+ RoutingMethodType.Default,
1390
+ RoutingMethodType.Renormalize,
1391
+ RoutingMethodType.RenormalizeNaive,
1392
+ ]
1393
+
1394
+ @staticmethod
1395
+ def _supports_router_logits_dtype(
1396
+ router_logits_dtype: torch.dtype | None,
1397
+ routing_method: RoutingMethodType,
1398
+ ) -> bool:
1399
+ return True
1400
+
1401
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1402
+ w13 = cpu_prepack_moe_weight_int8(layer.w13_weight, "vsx")
1403
+ w2 = cpu_prepack_moe_weight_int8(layer.w2_weight, "vsx")
1404
+ replace_parameter(layer, "w13_weight", w13)
1405
+ replace_parameter(layer, "w2_weight", w2)
1406
+
1407
+ def workspace_shapes(
1408
+ self,
1409
+ M: int,
1410
+ N: int,
1411
+ K: int,
1412
+ topk: int,
1413
+ global_num_experts: int,
1414
+ local_num_experts: int,
1415
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
1416
+ activation: MoEActivation,
1417
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
1418
+ # cpu_fused_moe_int8 manages its own scratch space.
1419
+ return (0,), (0,), (M, K)
1420
+
1421
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
1422
+ return TopKWeightAndReduceNoOP()
1423
+
1424
+ def apply(
1425
+ self,
1426
+ output: torch.Tensor,
1427
+ hidden_states: torch.Tensor,
1428
+ w1: torch.Tensor,
1429
+ w2: torch.Tensor,
1430
+ topk_weights: torch.Tensor,
1431
+ topk_ids: torch.Tensor,
1432
+ activation: MoEActivation,
1433
+ global_num_experts: int,
1434
+ expert_map: torch.Tensor | None,
1435
+ a1q_scale: torch.Tensor | None,
1436
+ a2_scale: torch.Tensor | None,
1437
+ workspace13: torch.Tensor,
1438
+ workspace2: torch.Tensor,
1439
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
1440
+ apply_router_weight_on_input: bool,
1441
+ ) -> None:
1442
+ # The modular prepare step applies router weights to the input when
1443
+ # apply_router_weight_on_input is enabled.
1444
+ assert self.w1_scale is not None
1445
+ assert self.w2_scale is not None
1446
+ cpu_fused_moe_int8(
1447
+ output,
1448
+ hidden_states,
1449
+ w1,
1450
+ w2,
1451
+ self.w1_scale,
1452
+ self.w2_scale,
1453
+ self.w1_bias,
1454
+ self.w2_bias,
1455
+ topk_weights,
1456
+ topk_ids,
1457
+ activation.value,
1458
+ "vsx",
1459
+ skip_weighted=apply_router_weight_on_input,
1460
+ )
1461
+
1462
+
1463
class ZenCPUExpertsInt8(mk.FusedMoEExpertsModular):
1464
"""AMD Zen INT8 MoE with per-token activation and channelwise weight
1465
quantization, dispatched through zentorch."""
vllm/model_executor/layers/fused_moe/oracle/int8.py
+7
-1
@@ -71,10 +71,16 @@ def backend_to_kernel_cls(
71
from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
72
ArmCPUExpertsInt8,
73
CPUExpertsInt8,
74
+ PowerCPUExpertsInt8,
75
ZenCPUExpertsInt8,
76
)
77
77
- return [ZenCPUExpertsInt8, ArmCPUExpertsInt8, CPUExpertsInt8]
78
+ return [
79
+ ZenCPUExpertsInt8,
80
+ PowerCPUExpertsInt8,
81
+ ArmCPUExpertsInt8,
82
+ CPUExpertsInt8,
83
+ ]
84
else:
85
raise ValueError(f"Unknown Int8 MoE backend: {backend.value}")
86