[Hardware][Power] Unqualized MoE Backend for Power (VSX) (#51624)
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>
Akash kaothalkar committed
Aug 13, 2026 at 08:26 UTC
61826c1c6fa301f1b0dc56156fd2f3fad5dfd138
6 files changed
+500
-6
csrc/cpu/cpu_fused_moe.cpp
+12
@@ -27,10 +27,22 @@
27
#define NEON_DISPATCH(...) case cpu_utils::ISA::NEON:
28
#endif
29
30
+#ifdef __powerpc64__
31
+ #include "cpu/micro_gemm/cpu_micro_gemm_vsx.hpp"
32
+ #define VSX_DISPATCH(...) \
33
+ case cpu_utils::ISA::VSX: { \
34
+ using gemm_t = cpu_micro_gemm::MicroGemm<cpu_utils::ISA::VSX, scalar_t>; \
35
+ return __VA_ARGS__(); \
36
+ }
37
+#else
38
+ #define VSX_DISPATCH(...) case cpu_utils::ISA::VSX:
39
+#endif
40
+
41
#define CPU_ISA_DISPATCH_IMPL(ISA_TYPE, ...) \
42
[&] { \
43
switch (ISA_TYPE) { \
44
AMX_DISPATCH(__VA_ARGS__) \
45
+ VSX_DISPATCH(__VA_ARGS__) \
46
case cpu_utils::ISA::VEC: { \
47
using gemm_t = \
48
cpu_micro_gemm::MicroGemm<cpu_utils::ISA::VEC, scalar_t>; \
csrc/cpu/micro_gemm/cpu_micro_gemm_vsx.hpp
new
+445
@@ -0,0 +1,445 @@
1
+#ifndef CPU_MICRO_GEMM_VSX_HPP
2
+#define CPU_MICRO_GEMM_VSX_HPP
3
+#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
4
+#include <altivec.h>
5
+
6
+namespace cpu_micro_gemm {
7
+namespace {
8
+
9
+// 8-16-16 pattern, 2 regs for A, 4 regs for B, 8 regs for C, [8, K] @ [K, 16]
10
+template <typename scalar_t>
11
+class TileGemmVSX {
12
+ public:
13
+ FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
14
+ switch (m) {
15
+ case 1:
16
+ gemm_micro<1>(CPU_MICRO_GEMM_PARAMS);
17
+ break;
18
+ case 2:
19
+ gemm_micro<2>(CPU_MICRO_GEMM_PARAMS);
20
+ break;
21
+ case 3:
22
+ gemm_micro<3>(CPU_MICRO_GEMM_PARAMS);
23
+ break;
24
+ case 4:
25
+ gemm_micro<4>(CPU_MICRO_GEMM_PARAMS);
26
+ break;
27
+ case 5:
28
+ gemm_micro<5>(CPU_MICRO_GEMM_PARAMS);
29
+ break;
30
+ case 6:
31
+ gemm_micro<6>(CPU_MICRO_GEMM_PARAMS);
32
+ break;
33
+ case 7:
34
+ gemm_micro<7>(CPU_MICRO_GEMM_PARAMS);
35
+ break;
36
+ case 8:
37
+ gemm_micro<8>(CPU_MICRO_GEMM_PARAMS);
38
+ break;
39
+ }
40
+ }
41
+
42
+ template <int32_t M>
43
+ static void gemm_micro_vsx_fallback(DEFINE_CPU_MICRO_GEMM_PARAMS) {
44
+ __vector float c_regs[M][4];
45
+
46
+ if (accum_c) {
47
+ for (int i = 0; i < M; i++) {
48
+ for (int j = 0; j < 4; j++) {
49
+ c_regs[i][j] = (__vector float)vec_xl(0, &c_ptr[i * ldc + j * 4]);
50
+ }
51
+ }
52
+ } else {
53
+ for (int i = 0; i < M; i++) {
54
+ for (int j = 0; j < 4; j++) {
55
+ c_regs[i][j] = (__vector float){0.0f, 0.0f, 0.0f, 0.0f};
56
+ }
57
+ }
58
+ }
59
+
60
+ const __vector unsigned short vzero = {0};
61
+ const __vector unsigned char mask_k0 = {0, 1, 2, 3, 8, 9, 10, 11,
62
+ 16, 17, 18, 19, 24, 25, 26, 27};
63
+ const __vector unsigned char mask_k1 = {4, 5, 6, 7, 12, 13, 14, 15,
64
+ 20, 21, 22, 23, 28, 29, 30, 31};
65
+
66
+ for (int32_t k_idx = 0; k_idx < k; k_idx += 2) {
67
+ __vector float B_k0[4];
68
+ __vector float B_k1[4];
69
+
70
+ for (int j = 0; j < 4; j++) {
71
+ __vector unsigned short vB_short = (__vector unsigned short)vec_xl(
72
+ 0, (const unsigned char*)&b_ptr[k_idx * 16 + j * 8]);
73
+ __vector unsigned int vB_f32_01 =
74
+ (__vector unsigned int)vec_mergeh(vzero, vB_short);
75
+ __vector unsigned int vB_f32_23 =
76
+ (__vector unsigned int)vec_mergel(vzero, vB_short);
77
+
78
+ B_k0[j] = vec_perm((__vector float)vB_f32_01, (__vector float)vB_f32_23,
79
+ mask_k0);
80
+ B_k1[j] = vec_perm((__vector float)vB_f32_01, (__vector float)vB_f32_23,
81
+ mask_k1);
82
+ }
83
+
84
+ for (int i = 0; i < M; i++) {
85
+ uint32_t a_val;
86
+ std::memcpy(&a_val, &a_ptr[i * lda + k_idx], 4);
87
+ uint32_t a0_u = (a_val & 0xFFFF) << 16;
88
+ uint32_t a1_u = a_val & 0xFFFF0000;
89
+ float a0_f, a1_f;
90
+ std::memcpy(&a0_f, &a0_u, 4);
91
+ std::memcpy(&a1_f, &a1_u, 4);
92
+
93
+ __vector float vA0 = vec_splats(a0_f);
94
+ __vector float vA1 = vec_splats(a1_f);
95
+
96
+#pragma GCC unroll 4
97
+ for (int j = 0; j < 4; j++) {
98
+ c_regs[i][j] = vec_madd(vA0, B_k0[j], c_regs[i][j]);
99
+ c_regs[i][j] = vec_madd(vA1, B_k1[j], c_regs[i][j]);
100
+ }
101
+ }
102
+ }
103
+
104
+ for (int i = 0; i < M; i++) {
105
+ for (int j = 0; j < 4; j++) {
106
+ vec_xst(c_regs[i][j], 0, &c_ptr[i * ldc + j * 4]);
107
+ }
108
+ }
109
+ }
110
+
111
+ template <int32_t M>
112
+ static void gemm_micro(DEFINE_CPU_MICRO_GEMM_PARAMS) {
113
+ static_assert(0 < M && M <= 8);
114
+
115
+ if constexpr (M <= 4) {
116
+ gemm_micro_vsx_fallback<M>(CPU_MICRO_GEMM_PARAMS);
117
+ return;
118
+ }
119
+
120
+ // Calculate how many 4x4 tiles we need in M dimension
121
+ constexpr int tiles_m = (M + 3) / 4;
122
+ constexpr int tiles_n =
123
+ 4; // Since NSize=16, we always process 4 tiles in N dimension
124
+
125
+ __vector_quad acc[tiles_m][tiles_n];
126
+ for (int i = 0; i < tiles_m; i++) {
127
+ for (int j = 0; j < tiles_n; j++) {
128
+ __builtin_mma_xxsetaccz(&acc[i][j]);
129
+ }
130
+ }
131
+
132
+ // Load initial accumulators if accum_c is true
133
+ if (accum_c) {
134
+ typedef float v4sf __attribute__((vector_size(16)));
135
+ for (int i = 0; i < tiles_m; i++) {
136
+ for (int j = 0; j < tiles_n; j++) {
137
+ v4sf tmp[4];
138
+ for (int r = 0; r < 4; r++) {
139
+ int r_idx = i * 4 + r;
140
+ if (r_idx < M) {
141
+ tmp[r] = (v4sf)vec_xl(0, &c_ptr[r_idx * ldc + j * 4]);
142
+ } else {
143
+ tmp[r] = (v4sf){0.0f, 0.0f, 0.0f, 0.0f};
144
+ }
145
+ }
146
+ __builtin_mma_build_acc(&acc[i][j], (__vector unsigned char)tmp[0],
147
+ (__vector unsigned char)tmp[1],
148
+ (__vector unsigned char)tmp[2],
149
+ (__vector unsigned char)tmp[3]);
150
+ }
151
+ }
152
+ }
153
+
154
+ int32_t k_idx = 0;
155
+ const __vector unsigned short vzero = {0};
156
+
157
+ // Main unrolled loop processing 8 columns (4 k-iterations) at a time
158
+ for (; k_idx <= k - 8; k_idx += 8) {
159
+ __vector unsigned int vA_0[tiles_m];
160
+ __vector unsigned int vA_2[tiles_m];
161
+ __vector unsigned int vA_4[tiles_m];
162
+ __vector unsigned int vA_6[tiles_m];
163
+
164
+ if constexpr (M >= 1) {
165
+ __vector unsigned short v0;
166
+ if constexpr (M >= 1)
167
+ v0 = (__vector unsigned short)vec_xl(
168
+ 0, (const unsigned char*)&a_ptr[0 * lda + k_idx]);
169
+ else
170
+ v0 = vzero;
171
+ __vector unsigned short v1;
172
+ if constexpr (M >= 2)
173
+ v1 = (__vector unsigned short)vec_xl(
174
+ 0, (const unsigned char*)&a_ptr[1 * lda + k_idx]);
175
+ else
176
+ v1 = vzero;
177
+ __vector unsigned short v2;
178
+ if constexpr (M >= 3)
179
+ v2 = (__vector unsigned short)vec_xl(
180
+ 0, (const unsigned char*)&a_ptr[2 * lda + k_idx]);
181
+ else
182
+ v2 = vzero;
183
+ __vector unsigned short v3;
184
+ if constexpr (M >= 4)
185
+ v3 = (__vector unsigned short)vec_xl(
186
+ 0, (const unsigned char*)&a_ptr[3 * lda + k_idx]);
187
+ else
188
+ v3 = vzero;
189
+
190
+ __vector unsigned int u0 = (__vector unsigned int)v0;
191
+ __vector unsigned int u1 = (__vector unsigned int)v1;
192
+ __vector unsigned int u2 = (__vector unsigned int)v2;
193
+ __vector unsigned int u3 = (__vector unsigned int)v3;
194
+
195
+ __vector unsigned int h01 = vec_mergeh(u0, u1);
196
+ __vector unsigned int h23 = vec_mergeh(u2, u3);
197
+ __vector unsigned int l01 = vec_mergel(u0, u1);
198
+ __vector unsigned int l23 = vec_mergel(u2, u3);
199
+
200
+ vA_0[0] = (__vector unsigned int)vec_mergeh(
201
+ (__vector unsigned long long)h01, (__vector unsigned long long)h23);
202
+ vA_2[0] = (__vector unsigned int)vec_mergel(
203
+ (__vector unsigned long long)h01, (__vector unsigned long long)h23);
204
+ vA_4[0] = (__vector unsigned int)vec_mergeh(
205
+ (__vector unsigned long long)l01, (__vector unsigned long long)l23);
206
+ vA_6[0] = (__vector unsigned int)vec_mergel(
207
+ (__vector unsigned long long)l01, (__vector unsigned long long)l23);
208
+ }
209
+
210
+ if constexpr (M >= 5) {
211
+ __vector unsigned short v4;
212
+ if constexpr (M >= 5)
213
+ v4 = (__vector unsigned short)vec_xl(
214
+ 0, (const unsigned char*)&a_ptr[4 * lda + k_idx]);
215
+ else
216
+ v4 = vzero;
217
+ __vector unsigned short v5;
218
+ if constexpr (M >= 6)
219
+ v5 = (__vector unsigned short)vec_xl(
220
+ 0, (const unsigned char*)&a_ptr[5 * lda + k_idx]);
221
+ else
222
+ v5 = vzero;
223
+ __vector unsigned short v6;
224
+ if constexpr (M >= 7)
225
+ v6 = (__vector unsigned short)vec_xl(
226
+ 0, (const unsigned char*)&a_ptr[6 * lda + k_idx]);
227
+ else
228
+ v6 = vzero;
229
+ __vector unsigned short v7;
230
+ if constexpr (M >= 8)
231
+ v7 = (__vector unsigned short)vec_xl(
232
+ 0, (const unsigned char*)&a_ptr[7 * lda + k_idx]);
233
+ else
234
+ v7 = vzero;
235
+
236
+ __vector unsigned int u4 = (__vector unsigned int)v4;
237
+ __vector unsigned int u5 = (__vector unsigned int)v5;
238
+ __vector unsigned int u6 = (__vector unsigned int)v6;
239
+ __vector unsigned int u7 = (__vector unsigned int)v7;
240
+
241
+ __vector unsigned int h45 = vec_mergeh(u4, u5);
242
+ __vector unsigned int h67 = vec_mergeh(u6, u7);
243
+ __vector unsigned int l45 = vec_mergel(u4, u5);
244
+ __vector unsigned int l67 = vec_mergel(u6, u7);
245
+
246
+ vA_0[1] = (__vector unsigned int)vec_mergeh(
247
+ (__vector unsigned long long)h45, (__vector unsigned long long)h67);
248
+ vA_2[1] = (__vector unsigned int)vec_mergel(
249
+ (__vector unsigned long long)h45, (__vector unsigned long long)h67);
250
+ vA_4[1] = (__vector unsigned int)vec_mergeh(
251
+ (__vector unsigned long long)l45, (__vector unsigned long long)l67);
252
+ vA_6[1] = (__vector unsigned int)vec_mergel(
253
+ (__vector unsigned long long)l45, (__vector unsigned long long)l67);
254
+ }
255
+
256
+ // Load B and GER for k_idx + 0
257
+ __vector unsigned char vB_vec_0[tiles_n];
258
+ for (int j = 0; j < tiles_n; j++)
259
+ vB_vec_0[j] =
260
+ vec_xl(0, (const unsigned char*)&b_ptr[(k_idx + 0) * 16 + j * 8]);
261
+ for (int i = 0; i < tiles_m; i++)
262
+ for (int j = 0; j < tiles_n; j++)
263
+ __builtin_mma_xvbf16ger2pp(
264
+ &acc[i][j], (__vector unsigned char)vA_0[i], vB_vec_0[j]);
265
+
266
+ // Load B and GER for k_idx + 2
267
+ __vector unsigned char vB_vec_2[tiles_n];
268
+ for (int j = 0; j < tiles_n; j++)
269
+ vB_vec_2[j] =
270
+ vec_xl(0, (const unsigned char*)&b_ptr[(k_idx + 2) * 16 + j * 8]);
271
+ for (int i = 0; i < tiles_m; i++)
272
+ for (int j = 0; j < tiles_n; j++)
273
+ __builtin_mma_xvbf16ger2pp(
274
+ &acc[i][j], (__vector unsigned char)vA_2[i], vB_vec_2[j]);
275
+
276
+ // Load B and GER for k_idx + 4
277
+ __vector unsigned char vB_vec_4[tiles_n];
278
+ for (int j = 0; j < tiles_n; j++)
279
+ vB_vec_4[j] =
280
+ vec_xl(0, (const unsigned char*)&b_ptr[(k_idx + 4) * 16 + j * 8]);
281
+ for (int i = 0; i < tiles_m; i++)
282
+ for (int j = 0; j < tiles_n; j++)
283
+ __builtin_mma_xvbf16ger2pp(
284
+ &acc[i][j], (__vector unsigned char)vA_4[i], vB_vec_4[j]);
285
+
286
+ // Load B and GER for k_idx + 6
287
+ __vector unsigned char vB_vec_6[tiles_n];
288
+ for (int j = 0; j < tiles_n; j++)
289
+ vB_vec_6[j] =
290
+ vec_xl(0, (const unsigned char*)&b_ptr[(k_idx + 6) * 16 + j * 8]);
291
+ for (int i = 0; i < tiles_m; i++)
292
+ for (int j = 0; j < tiles_n; j++)
293
+ __builtin_mma_xvbf16ger2pp(
294
+ &acc[i][j], (__vector unsigned char)vA_6[i], vB_vec_6[j]);
295
+ }
296
+
297
+ // Remainder loop processing 2 columns at a time
298
+ for (; k_idx < k; k_idx += 2) {
299
+ __vector unsigned int vA_uint[tiles_m];
300
+ if constexpr (M == 1) {
301
+ uint32_t val0;
302
+ std::memcpy(&val0, &a_ptr[0 * lda + k_idx], 4);
303
+ vA_uint[0] = (__vector unsigned int){val0, 0, 0, 0};
304
+ } else if constexpr (M == 2) {
305
+ uint32_t val0, val1;
306
+ std::memcpy(&val0, &a_ptr[0 * lda + k_idx], 4);
307
+ std::memcpy(&val1, &a_ptr[1 * lda + k_idx], 4);
308
+ vA_uint[0] = (__vector unsigned int){val0, val1, 0, 0};
309
+ } else if constexpr (M == 3) {
310
+ uint32_t val0, val1, val2;
311
+ std::memcpy(&val0, &a_ptr[0 * lda + k_idx], 4);
312
+ std::memcpy(&val1, &a_ptr[1 * lda + k_idx], 4);
313
+ std::memcpy(&val2, &a_ptr[2 * lda + k_idx], 4);
314
+ vA_uint[0] = (__vector unsigned int){val0, val1, val2, 0};
315
+ } else if constexpr (M == 4) {
316
+ uint32_t val0, val1, val2, val3;
317
+ std::memcpy(&val0, &a_ptr[0 * lda + k_idx], 4);
318
+ std::memcpy(&val1, &a_ptr[1 * lda + k_idx], 4);
319
+ std::memcpy(&val2, &a_ptr[2 * lda + k_idx], 4);
320
+ std::memcpy(&val3, &a_ptr[3 * lda + k_idx], 4);
321
+ vA_uint[0] = (__vector unsigned int){val0, val1, val2, val3};
322
+ } else if constexpr (M == 5) {
323
+ uint32_t val0, val1, val2, val3, val4;
324
+ std::memcpy(&val0, &a_ptr[0 * lda + k_idx], 4);
325
+ std::memcpy(&val1, &a_ptr[1 * lda + k_idx], 4);
326
+ std::memcpy(&val2, &a_ptr[2 * lda + k_idx], 4);
327
+ std::memcpy(&val3, &a_ptr[3 * lda + k_idx], 4);
328
+ vA_uint[0] = (__vector unsigned int){val0, val1, val2, val3};
329
+ std::memcpy(&val4, &a_ptr[4 * lda + k_idx], 4);
330
+ vA_uint[1] = (__vector unsigned int){val4, 0, 0, 0};
331
+ } else if constexpr (M == 6) {
332
+ uint32_t val0, val1, val2, val3, val4, val5;
333
+ std::memcpy(&val0, &a_ptr[0 * lda + k_idx], 4);
334
+ std::memcpy(&val1, &a_ptr[1 * lda + k_idx], 4);
335
+ std::memcpy(&val2, &a_ptr[2 * lda + k_idx], 4);
336
+ std::memcpy(&val3, &a_ptr[3 * lda + k_idx], 4);
337
+ vA_uint[0] = (__vector unsigned int){val0, val1, val2, val3};
338
+ std::memcpy(&val4, &a_ptr[4 * lda + k_idx], 4);
339
+ std::memcpy(&val5, &a_ptr[5 * lda + k_idx], 4);
340
+ vA_uint[1] = (__vector unsigned int){val4, val5, 0, 0};
341
+ } else if constexpr (M == 7) {
342
+ uint32_t val0, val1, val2, val3, val4, val5, val6;
343
+ std::memcpy(&val0, &a_ptr[0 * lda + k_idx], 4);
344
+ std::memcpy(&val1, &a_ptr[1 * lda + k_idx], 4);
345
+ std::memcpy(&val2, &a_ptr[2 * lda + k_idx], 4);
346
+ std::memcpy(&val3, &a_ptr[3 * lda + k_idx], 4);
347
+ vA_uint[0] = (__vector unsigned int){val0, val1, val2, val3};
348
+ std::memcpy(&val4, &a_ptr[4 * lda + k_idx], 4);
349
+ std::memcpy(&val5, &a_ptr[5 * lda + k_idx], 4);
350
+ std::memcpy(&val6, &a_ptr[6 * lda + k_idx], 4);
351
+ vA_uint[1] = (__vector unsigned int){val4, val5, val6, 0};
352
+ } else if constexpr (M == 8) {
353
+ uint32_t val0, val1, val2, val3, val4, val5, val6, val7;
354
+ std::memcpy(&val0, &a_ptr[0 * lda + k_idx], 4);
355
+ std::memcpy(&val1, &a_ptr[1 * lda + k_idx], 4);
356
+ std::memcpy(&val2, &a_ptr[2 * lda + k_idx], 4);
357
+ std::memcpy(&val3, &a_ptr[3 * lda + k_idx], 4);
358
+ vA_uint[0] = (__vector unsigned int){val0, val1, val2, val3};
359
+ std::memcpy(&val4, &a_ptr[4 * lda + k_idx], 4);
360
+ std::memcpy(&val5, &a_ptr[5 * lda + k_idx], 4);
361
+ std::memcpy(&val6, &a_ptr[6 * lda + k_idx], 4);
362
+ std::memcpy(&val7, &a_ptr[7 * lda + k_idx], 4);
363
+ vA_uint[1] = (__vector unsigned int){val4, val5, val6, val7};
364
+ }
365
+
366
+ // Load packed B
367
+ __vector unsigned char vB_vec[tiles_n];
368
+ for (int j = 0; j < tiles_n; j++) {
369
+ vB_vec[j] = vec_xl(0, (const unsigned char*)&b_ptr[k_idx * 16 + j * 8]);
370
+ }
371
+
372
+ for (int i = 0; i < tiles_m; i++) {
373
+ for (int j = 0; j < tiles_n; j++) {
374
+ __builtin_mma_xvbf16ger2pp(
375
+ &acc[i][j], (__vector unsigned char)vA_uint[i], vB_vec[j]);
376
+ }
377
+ }
378
+ }
379
+
380
+ // Disassemble and store
381
+ typedef float v4sf __attribute__((vector_size(16)));
382
+ for (int i = 0; i < tiles_m; i++) {
383
+ for (int j = 0; j < tiles_n; j++) {
384
+ v4sf tmp[4];
385
+ __builtin_mma_disassemble_acc(tmp, &acc[i][j]);
386
+ for (int r = 0; r < 4; r++) {
387
+ int r_idx = i * 4 + r;
388
+ if (r_idx < M) {
389
+ vec_xst((__vector float)tmp[r], 0, &c_ptr[r_idx * ldc + j * 4]);
390
+ }
391
+ }
392
+ }
393
+ }
394
+ }
395
+};
396
+} // namespace
397
+
398
+// Gemm kernel uses MMA instructions, requires B matrix to be packed
399
+template <typename scalar_t>
400
+class MicroGemm<cpu_utils::ISA::VSX, scalar_t> {
401
+ public:
402
+ static constexpr int32_t MaxMSize = 8;
403
+ static constexpr int32_t NSize = 16;
404
+ static constexpr int32_t WeightOCGroupSize = 16;
405
+ static constexpr bool PackA = false;
406
+
407
+ public:
408
+ void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
409
+ TileGemmVSX<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
410
+ }
411
+
412
+ // Pack weight:
413
+ // Original B is (output_size, input_size).
414
+ // We need to pack it for MMA. MMA `xvbf16ger2pp` requires 4 output columns
415
+ // and 2 input columns in a single VecBF16. We pack into shape: [output_size /
416
+ // 4, input_size / 2, 8]
417
+ static void pack_weight(const scalar_t* __restrict__ weight,
418
+ scalar_t* __restrict__ packed_weight,
419
+ const int32_t output_size, const int32_t input_size) {
420
+ TORCH_CHECK_EQ(output_size % NSize, 0);
421
+ TORCH_CHECK_EQ(input_size % 2, 0);
422
+
423
+ scalar_t* pw = packed_weight;
424
+ for (int32_t o_block = 0; o_block < output_size; o_block += NSize) {
425
+ for (int32_t i_idx = 0; i_idx < input_size; i_idx += 2) {
426
+ for (int32_t o_idx = o_block; o_idx < o_block + NSize; o_idx += 4) {
427
+ *pw++ = weight[(o_idx + 0) * input_size + i_idx];
428
+ *pw++ = weight[(o_idx + 0) * input_size + i_idx + 1];
429
+
430
+ *pw++ = weight[(o_idx + 1) * input_size + i_idx];
431
+ *pw++ = weight[(o_idx + 1) * input_size + i_idx + 1];
432
+
433
+ *pw++ = weight[(o_idx + 2) * input_size + i_idx];
434
+ *pw++ = weight[(o_idx + 2) * input_size + i_idx + 1];
435
+
436
+ *pw++ = weight[(o_idx + 3) * input_size + i_idx];
437
+ *pw++ = weight[(o_idx + 3) * input_size + i_idx + 1];
438
+ }
439
+ }
440
+ }
441
+ }
442
+};
443
+} // namespace cpu_micro_gemm
444
+
445
+#endif
csrc/cpu/utils.hpp
+3
-1
@@ -9,7 +9,7 @@
9
#include "cpu/cpu_types.hpp"
10
11
namespace cpu_utils {
12
-enum class ISA { AMX, VEC, RVV, NEON };
12
+enum class ISA { AMX, VEC, RVV, NEON, VSX };
13
14
inline ISA get_isa(const std::string& isa) {
15
if (isa == "amx") {
@@ -20,6 +20,8 @@ inline ISA get_isa(const std::string& isa) {
20
return ISA::RVV;
21
} else if (isa == "neon") {
22
return ISA::NEON;
23
+ } else if (isa == "vsx") {
24
+ return ISA::VSX;
25
} else {
26
TORCH_CHECK(false, "Invalid isa type: " + isa);
27
}
tests/kernels/moe/test_cpu_fused_moe.py
+5
-5
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import sys
4
5
import pytest
6
import torch
@@ -25,6 +24,7 @@ from vllm.model_executor.layers.fused_moe.config import (
24
from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
25
ArmCPUUnquantizedExperts,
26
CPUUnquantizedExperts,
27
+ PowerCPUUnquantizedExperts,
28
X86CPUUnquantizedExperts,
29
select_experts,
30
)
@@ -48,13 +48,12 @@ ACT = [
48
]
49
USE_BIAS = [False, True]
50
ISA = ["vec"]
51
-if (
52
- current_platform.get_cpu_architecture() == CpuArchEnum.ARM
53
- and sys.platform != "darwin"
54
-):
51
+if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
52
ISA.append("neon")
53
if torch.cpu._is_amx_tile_supported():
54
ISA.append("amx")
55
+if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC:
56
+ ISA.append("vsx")
57
58
DTYPE = [torch.bfloat16]
59
@@ -551,6 +550,7 @@ def test_cpu_fused_moe_unaligned_intermediate_size(
550
"vec": CPUUnquantizedExperts,
551
"amx": X86CPUUnquantizedExperts,
552
"neon": ArmCPUUnquantizedExperts,
553
+ "vsx": PowerCPUUnquantizedExperts,
554
}[isa]
555
supported, reason = experts_cls.is_supported_config(
556
experts_cls,
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+33
@@ -500,6 +500,39 @@ class ArmCPUUnquantizedExperts(CPUUnquantizedExperts):
500
return cpu_cls._supports_grouped_gemm(moe_config)
501
502
503
+class PowerCPUUnquantizedExperts(CPUUnquantizedExperts):
504
+ """PowerPC VSX grouped-gemm unquantized MoE experts."""
505
+
506
+ isa = "vsx"
507
+ output_alignment = 16
508
+ reduction_alignment = 2
509
+
510
+ @staticmethod
511
+ def _supports_current_device() -> bool:
512
+ return (
513
+ current_platform.is_cpu()
514
+ and current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC
515
+ )
516
+
517
+ @staticmethod
518
+ def is_supported_config(
519
+ cls: type[mk.FusedMoEExperts],
520
+ moe_config: FusedMoEConfig,
521
+ weight_key: QuantKey | None,
522
+ activation_key: QuantKey | None,
523
+ activation_format: mk.FusedMoEActivationFormat,
524
+ ) -> tuple[bool, str | None]:
525
+ supported, reason = mk.FusedMoEExperts.is_supported_config(
526
+ cls, moe_config, weight_key, activation_key, activation_format
527
+ )
528
+ if not supported:
529
+ return supported, reason
530
+ if moe_config.in_dtype != torch.bfloat16:
531
+ return False, "kernel requires bfloat16 activations"
532
+ cpu_cls = cast(type[CPUUnquantizedExperts], cls)
533
+ return cpu_cls._supports_grouped_gemm(moe_config)
534
+
535
+
536
# ===========================================================================
537
# FP8 W8A16 MoE
538
# ===========================================================================
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+2
@@ -147,6 +147,7 @@ def backend_to_kernel_cls(
147
from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
148
ArmCPUUnquantizedExperts,
149
CPUUnquantizedExperts,
150
+ PowerCPUUnquantizedExperts,
151
X86CPUUnquantizedExperts,
152
)
153
@@ -154,6 +155,7 @@ def backend_to_kernel_cls(
155
return [
156
X86CPUUnquantizedExperts,
157
ArmCPUUnquantizedExperts,
158
+ PowerCPUUnquantizedExperts,
159
CPUUnquantizedExperts,
160
]
161