@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 420b0a5c95
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 07:30 UTC

[Hardware][Power]Add Power VSX Attention Backend and fix l2 Cache Crash (#40451)

Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com> Signed-off-by: Akash Kaothalkar <akash.kaothalkar@ibm.com> Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com> Co-authored-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com> Co-authored-by: Akash Kaothalkar <akash.kaothalkar@ibm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>

Akash kaothalkar committed May 5, 2026 at 09:21 UTC 420b0a5c95187809b2701323f1112472b2f3b707
8 files changed +399 -9
csrc/cpu/cpu_attn.cpp
+4
@@ -29,6 +29,8 @@ torch::Tensor get_scheduler_metadata(
29 isa = cpu_attention::ISA::NEON;
30 } else if (isa_hint == "vxe") {
31 isa = cpu_attention::ISA::VXE;
32 + } else if (isa_hint == "vsx") {
33 + isa = cpu_attention::ISA::VSX;
34 } else {
35 TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
36 }
@@ -129,6 +131,8 @@ void cpu_attn_reshape_and_cache(
131 return cpu_attention::ISA::NEON;
132 } else if (isa == "vxe") {
133 return cpu_attention::ISA::VXE;
134 + } else if (isa == "vsx") {
135 + return cpu_attention::ISA::VSX;
136 } else {
137 TORCH_CHECK(false, "Invalid ISA type: " + isa);
138 }
csrc/cpu/cpu_attn_impl.hpp
+4 -1
@@ -12,7 +12,7 @@
12 #include "cpu/utils.hpp"
13
14 namespace cpu_attention {
15 -enum class ISA { AMX, VEC, VEC16, NEON, VXE };
15 +enum class ISA { AMX, VEC, VEC16, NEON, VXE, VSX };
16
17 // Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.
18 enum class Fp8KVCacheDataType {
@@ -164,6 +164,9 @@ struct AttentionMetadata {
164 case ISA::VXE:
165 ss << "VXE, ";
166 break;
167 + case ISA::VSX:
168 + ss << "VSX, ";
169 + break;
170 }
171 ss << "workitem_group_num: " << workitem_group_num
172 << ", reduction_item_num: " << reduction_item_num
csrc/cpu/cpu_attn_vec.hpp
+2 -2
@@ -27,8 +27,8 @@ FORCE_INLINE std::pair<vec_op::FP32Vec16, vec_op::FP32Vec16> load_b_pair_vec(
27 return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
28 } else {
29 using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
30 - return {vec_op::FP32Vec16(load_vec_t(ptr)),
31 - vec_op::FP32Vec16(load_vec_t(ptr + 16))};
30 + return std::make_pair(vec_op::FP32Vec16(load_vec_t(ptr)),
31 + vec_op::FP32Vec16(load_vec_t(ptr + 16)));
32 }
33 }
34
csrc/cpu/cpu_attn_vsx.hpp new
+359
@@ -0,0 +1,359 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +#ifndef CPU_ATTN_VSX_HPP
4 +#define CPU_ATTN_VSX_HPP
5 +
6 +#include "cpu_attn_impl.hpp"
7 +#include <altivec.h>
8 +#include <type_traits>
9 +
10 +namespace cpu_attention {
11 +
12 +namespace {
13 +
14 +// ppc64le Vector = 16 bytes (128 bits)
15 +#define BLOCK_SIZE_ALIGNMENT 32
16 +#define HEAD_SIZE_ALIGNMENT 32
17 +#define MAX_Q_HEAD_NUM_PER_ITER 16
18 +
19 +template <typename kv_cache_t>
20 +FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, __vector float& b0,
21 + __vector float& b1);
22 +
23 +// [1] Float Specialization
24 +template <>
25 +FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, __vector float& b0,
26 + __vector float& b1) {
27 + b0 = vec_xl(0, const_cast<float*>(p));
28 + b1 = vec_xl(0, const_cast<float*>(p + 4));
29 +}
30 +
31 +// [2] BFloat16 Specialization (Little Endian ppc64le)
32 +// On ppc64le (LE): BF16 bits should land in the HIGH 16 bits of each float32.
33 +// Byte layout of float32 on LE: [byte0(LSB), byte1, byte2, byte3(MSB)]
34 +// We need BF16 in bytes2-3 (high half) with bytes0-1 zeroed.
35 +// vec_mergeh on LE interleaves elements 0..3: result_i = {a[i], b[i]}
36 +// So vec_mergeh(zeros_u16, raw_u16) gives for each uint16 pair:
37 +// uint16[2i] = zeros[i] -> low 16 bits of uint32 -> zeroed mantissa LSBs
38 +// uint16[2i+1] = raw[i] -> high 16 bits of uint32 -> BF16 bits
39 +// Cast to float32 gives exactly (bf16_bits << 16) per element.
40 +template <>
41 +FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
42 + __vector float& b0,
43 + __vector float& b1) {
44 + __vector unsigned short raw = vec_xl(
45 + 0, reinterpret_cast<unsigned short*>(const_cast<c10::BFloat16*>(p)));
46 + __vector unsigned short zeros = vec_splat_u16(0);
47 +
48 + // LE: zeros in low 16 bits, raw in high 16 bits → bf16 << 16 == float32
49 + b0 = (__vector float)vec_mergeh(zeros, raw);
50 + b1 = (__vector float)vec_mergel(zeros, raw);
51 +}
52 +
53 +// Note: c10::Half (FP16) is not supported on PowerPC architecture
54 +
55 +template <int32_t M, typename kv_cache_t>
56 +FORCE_INLINE void gemm_micro_ppc64le_Mx8_Ku4(
57 + const float* __restrict A, // [M x K]
58 + const kv_cache_t* __restrict B, // [K x 8]
59 + float* __restrict C, // [M x 8]
60 + int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
61 + static_assert(1 <= M && M <= 8, "M must be in [1,8]");
62 +
63 +#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
64 +#define IF_M(i) if constexpr (M > (i))
65 +
66 + // 1. Define A pointers
67 +#define DECL_A(i) const float* a##i = A + (i) * lda;
68 + ROWS_APPLY(DECL_A)
69 +#undef DECL_A
70 +
71 + // 2. Define Accumulators (2 vectors covers 8 columns)
72 +#define DECL_ACC(i) __vector float acc##i##_0, acc##i##_1;
73 + ROWS_APPLY(DECL_ACC)
74 +#undef DECL_ACC
75 +
76 + // 3. Initialize Accumulators (Load C or Zero)
77 +#define INIT_ACC(i) \
78 + IF_M(i) { \
79 + if (accumulate) { \
80 + acc##i##_0 = vec_xl(0, const_cast<float*>(C + (i) * ldc + 0)); \
81 + acc##i##_1 = vec_xl(0, const_cast<float*>(C + (i) * ldc + 4)); \
82 + } else { \
83 + acc##i##_0 = vec_splats(0.0f); \
84 + acc##i##_1 = vec_splats(0.0f); \
85 + } \
86 + }
87 + ROWS_APPLY(INIT_ACC)
88 +#undef INIT_ACC
89 +
90 + int32_t k = 0;
91 +
92 + for (; k + 3 < K; k += 4) {
93 + // Load 4 values of A for each Row M: A[k...k+3]
94 +#define LOAD_A4(i) \
95 + __vector float a##i##v; \
96 + IF_M(i) a##i##v = vec_xl(0, const_cast<float*>(a##i + k));
97 + ROWS_APPLY(LOAD_A4)
98 +#undef LOAD_A4
99 +
100 + // FMA for specific lane L of A
101 + // ppc64le: vec_madd(b, vec_splat(a, lane), acc)
102 +#define FMAS_LANE(i, aiv, L) \
103 + IF_M(i) { \
104 + __vector float a_broad = vec_splat(aiv, L); \
105 + acc##i##_0 = vec_madd(b0, a_broad, acc##i##_0); \
106 + acc##i##_1 = vec_madd(b1, a_broad, acc##i##_1); \
107 + }
108 +
109 + // Unroll K=0..3
110 + {
111 + __vector float b0, b1;
112 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1);
113 +#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
114 + ROWS_APPLY(STEP_K0)
115 +#undef STEP_K0
116 + }
117 + {
118 + __vector float b0, b1;
119 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1);
120 +#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
121 + ROWS_APPLY(STEP_K1)
122 +#undef STEP_K1
123 + }
124 + {
125 + __vector float b0, b1;
126 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1);
127 +#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
128 + ROWS_APPLY(STEP_K2)
129 +#undef STEP_K2
130 + }
131 + {
132 + __vector float b0, b1;
133 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1);
134 +#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
135 + ROWS_APPLY(STEP_K3)
136 +#undef STEP_K3
137 + }
138 +#undef FMAS_LANE
139 + }
140 +
141 + for (; k < K; ++k) {
142 + __vector float b0, b1;
143 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
144 +#define TAIL_ROW(i) \
145 + IF_M(i) { \
146 + __vector float ai = vec_splats(*(a##i + k)); \
147 + acc##i##_0 = vec_madd(b0, ai, acc##i##_0); \
148 + acc##i##_1 = vec_madd(b1, ai, acc##i##_1); \
149 + }
150 + ROWS_APPLY(TAIL_ROW)
151 +#undef TAIL_ROW
152 + }
153 +
154 +#define STORE_ROW(i) \
155 + IF_M(i) { \
156 + vec_xst(acc##i##_0, 0, C + (i) * ldc + 0); \
157 + vec_xst(acc##i##_1, 0, C + (i) * ldc + 4); \
158 + }
159 + ROWS_APPLY(STORE_ROW)
160 +#undef STORE_ROW
161 +
162 +#undef ROWS_APPLY
163 +#undef IF_M
164 +}
165 +
166 +template <int32_t N, typename kv_cache_t>
167 +FORCE_INLINE void gemm_macro_ppc64le_Mx8_Ku4(const float* __restrict A,
168 + const kv_cache_t* __restrict B,
169 + float* __restrict C, int32_t M,
170 + int32_t K, int64_t lda,
171 + int64_t ldb, int64_t ldc,
172 + bool accumulate) {
173 + static_assert(N % 8 == 0, "N must be a multiple of 8");
174 + for (int32_t m = 0; m < M;) {
175 + int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
176 + const float* Ab = A + m * lda;
177 + float* Cb = C + m * ldc;
178 +
179 + for (int32_t n = 0; n < N; n += 8) {
180 + const kv_cache_t* Bn = B + n;
181 + float* Cn = Cb + n;
182 + switch (mb) {
183 + case 8:
184 + gemm_micro_ppc64le_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
185 + K, accumulate);
186 + break;
187 + case 4:
188 + gemm_micro_ppc64le_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
189 + K, accumulate);
190 + break;
191 + case 2:
192 + gemm_micro_ppc64le_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
193 + K, accumulate);
194 + break;
195 + default:
196 + gemm_micro_ppc64le_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
197 + K, accumulate);
198 + break;
199 + }
200 + }
201 + m += mb;
202 + }
203 +}
204 +
205 +template <typename kv_cache_t>
206 +class TileGemmPPC64 {
207 + public:
208 + template <AttentionGemmPhase phase, int32_t k_size>
209 + FORCE_INLINE static void gemm(const int32_t m_size,
210 + float* __restrict__ a_tile,
211 + kv_cache_t* __restrict__ b_tile,
212 + float* __restrict__ c_tile, const int64_t lda,
213 + const int64_t ldb, const int64_t ldc,
214 + const int32_t block_size,
215 + const int32_t dynamic_k_size,
216 + const bool accum_c) {
217 + if constexpr (phase == AttentionGemmPhase::QK) {
218 + gemm_macro_ppc64le_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
219 + a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
220 + } else {
221 + gemm_macro_ppc64le_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
222 + a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
223 + accum_c);
224 + }
225 + }
226 +};
227 +
228 +} // namespace
229 +
230 +template <typename scalar_t, int64_t head_dim>
231 +class AttentionImpl<ISA::VSX, scalar_t, head_dim> {
232 + public:
233 + using query_t = scalar_t;
234 + using q_buffer_t = float;
235 + using kv_cache_t = scalar_t;
236 + using logits_buffer_t = float;
237 + using partial_output_buffer_t = float;
238 + using prob_buffer_t = float;
239 +
240 + constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT;
241 + constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT;
242 + constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
243 + constexpr static int64_t HeadDim = head_dim;
244 + constexpr static ISA ISAType = ISA::VSX;
245 + constexpr static bool scale_on_logits =
246 + false; // Scale is applied to Q during copy
247 +
248 + public:
249 + AttentionImpl() {}
250 +
251 + template <template <typename tile_gemm_t> typename attention>
252 + FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
253 + attention<TileGemmPPC64<kv_cache_t>> attention_iteration;
254 + attention_iteration(CPU_ATTENTION_PARAMS);
255 + }
256 +
257 + // Strides for Memory Layout
258 + constexpr static int64_t k_cache_token_group_stride(
259 + const int32_t block_size) {
260 + return BlockSizeAlignment; // [head_dim, block_size] layout
261 + }
262 +
263 + constexpr static int64_t v_cache_token_group_stride(
264 + const int32_t block_size) {
265 + return head_dim * BlockSizeAlignment;
266 + }
267 +
268 + constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
269 + return HeadDimAlignment;
270 + }
271 +
272 + static void copy_q_heads_tile(scalar_t* __restrict__ src,
273 + float* __restrict__ q_buffer,
274 + const int32_t q_num,
275 + const int32_t q_heads_per_kv,
276 + const int64_t q_num_stride,
277 + const int64_t q_head_stride, float scale) {
278 + __vector float scale_vec = vec_splats(scale);
279 + constexpr bool is_bf16 = std::is_same<scalar_t, c10::BFloat16>::value;
280 +
281 + for (int32_t i = 0; i < q_num; ++i) {
282 + for (int32_t h = 0; h < q_heads_per_kv; ++h) {
283 + scalar_t* curr_src = src + i * q_num_stride + h * q_head_stride;
284 + float* curr_dst =
285 + q_buffer + i * q_heads_per_kv * head_dim + h * head_dim;
286 +
287 + int32_t d = 0;
288 + for (; d <= head_dim - 8; d += 8) {
289 + __vector float v0, v1;
290 + load_row8_B_as_f32<scalar_t>(curr_src + d, v0, v1);
291 +
292 + v0 = vec_mul(v0, scale_vec);
293 + v1 = vec_mul(v1, scale_vec);
294 +
295 + vec_xst(v0, 0, curr_dst + d);
296 + vec_xst(v1, 0, curr_dst + d + 4);
297 + }
298 +
299 + for (; d < head_dim; ++d) {
300 + float val = static_cast<float>(curr_src[d]);
301 + curr_dst[d] = val * scale;
302 + }
303 + }
304 + }
305 + }
306 +
307 + static void reshape_and_cache(
308 + const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
309 + scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
310 + const int64_t* __restrict__ slot_mapping, const int64_t token_num,
311 + const int64_t key_token_num_stride, const int64_t value_token_num_stride,
312 + const int64_t head_num, const int64_t key_head_num_stride,
313 + const int64_t value_head_num_stride, const int64_t num_blocks,
314 + const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
315 + const int64_t block_size, const int64_t block_size_stride,
316 + const float k_inv = 0.0f, const float v_inv = 0.0f) {
317 + // k_inv and v_inv are unused on VSX: FP8 KV cache is not supported on
318 + // PowerPC. The parameters are present to match the common interface.
319 +#pragma omp parallel for collapse(2)
320 + for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
321 + for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
322 + const int64_t pos = slot_mapping[token_idx];
323 + if (pos < 0) continue;
324 +
325 + const int64_t block_idx = pos / block_size;
326 + const int64_t block_offset = pos % block_size;
327 +
328 + {
329 + const scalar_t* key_src = key + token_idx * key_token_num_stride +
330 + head_idx * key_head_num_stride;
331 + scalar_t* key_dst = key_cache + block_idx * num_blocks_stride +
332 + head_idx * cache_head_num_stride + block_offset;
333 +
334 + for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
335 + key_dst[j] = key_src[i];
336 + }
337 + }
338 +
339 + {
340 + const scalar_t* val_src = value + token_idx * value_token_num_stride +
341 + head_idx * value_head_num_stride;
342 + scalar_t* val_dst = value_cache + block_idx * num_blocks_stride +
343 + head_idx * cache_head_num_stride +
344 + block_offset * head_dim;
345 +
346 + std::memcpy(val_dst, val_src, sizeof(scalar_t) * head_dim);
347 + }
348 + }
349 + }
350 + }
351 +};
352 +
353 +} // namespace cpu_attention
354 +
355 +#undef BLOCK_SIZE_ALIGNMENT
356 +#undef HEAD_SIZE_ALIGNMENT
357 +#undef MAX_Q_HEAD_NUM_PER_ITER
358 +
359 +#endif // CPU_ATTN_VSX_HPP
csrc/cpu/cpu_types_vsx.hpp
+4
@@ -9,6 +9,10 @@
9
10 namespace vec_op {
11
12 +// FP8 tag types for tag dispatch (see cpu_attn_vec.hpp)
13 +struct fp8_e4m3_tag {};
14 +struct fp8_e5m2_tag {};
15 +
16 // FIXME: FP16 is not fully supported in Torch-CPU
17 #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
18 AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
csrc/cpu/generate_cpu_attn_dispatch.py
+13 -2
@@ -20,6 +20,7 @@ ISA_TYPES = {
20 "VEC16": 2,
21 "NEON": 3,
22 "VXE": 4,
23 + "VSX": 5,
24 }
25
26 # KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2
@@ -37,7 +38,7 @@ KV_CACHE_CPP_TYPES = {
38 }
39
40 # ISAs supported for head_dims divisible by 32
40 -ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE"]
41 +ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE", "VSX"]
42
43 # ISAs supported for head_dims divisible by 16 only
44 ISA_FOR_16 = ["VEC16"]
@@ -148,6 +149,10 @@ def generate_header_file() -> str:
149 #include "cpu_attn_vxe.hpp"
150 #endif
151
152 +#ifdef __powerpc__
153 + #include "cpu_attn_vsx.hpp"
154 +#endif
155 +
156 """
157
158 header += generate_helper_function()
@@ -207,6 +212,11 @@ def generate_header_file() -> str:
212 ["VXE", "VEC", "VEC16"],
213 fp8=False,
214 )
215 + header += _macro_block(
216 + "#elif defined(__powerpc__)",
217 + ["VSX", "VEC", "VEC16"],
218 + fp8=False,
219 + )
220 header += _macro_block(
221 "#elif defined(__AVX512F__)",
222 ["VEC", "VEC16"],
@@ -223,7 +233,8 @@ def generate_header_file() -> str:
233 fp8=False,
234 )
235 header += (
226 - "#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */\n\n"
236 + "#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / "
237 + "__s390x__ / __powerpc__ */\n\n"
238 "#endif // CPU_ATTN_DISPATCH_GENERATED_H\n"
239 )
240
csrc/cpu/utils.hpp
+1 -1
@@ -54,7 +54,7 @@ struct Counter {
54 };
55
56 inline int64_t get_available_l2_size() {
57 -#if defined(__s390x__)
57 +#if defined(__s390x__) || defined(__powerpc__)
58 static int64_t size = []() {
59 uint32_t l2_cache_size = 0;
60 auto caps = at::cpu::get_cpu_capabilities();
vllm/v1/attention/backends/cpu_attn.py
+12 -3
@@ -29,7 +29,12 @@ from vllm.v1.kv_cache_interface import AttentionSpec, CrossAttentionSpec
29
30 logger = init_logger(__name__)
31
32 -_CPU_ARCH_PREFER_MIXED_BATCH = (CpuArchEnum.X86, CpuArchEnum.ARM, CpuArchEnum.S390X)
32 +_CPU_ARCH_PREFER_MIXED_BATCH = (
33 + CpuArchEnum.X86,
34 + CpuArchEnum.ARM,
35 + CpuArchEnum.S390X,
36 + CpuArchEnum.POWERPC,
37 +)
38
39
40 class CPUAttentionBackend(AttentionBackend):
@@ -510,8 +515,10 @@ def _get_attn_isa(
515 )
516 return "vec16"
517 supports_amx = torch.cpu._is_amx_tile_supported()
513 - supports_arm = current_platform.get_cpu_architecture() == CpuArchEnum.ARM
514 - supports_vxe = current_platform.get_cpu_architecture() == CpuArchEnum.S390X
518 + arch = current_platform.get_cpu_architecture()
519 + supports_arm = arch == CpuArchEnum.ARM
520 + supports_vxe = arch == CpuArchEnum.S390X
521 + supports_vsx = arch == CpuArchEnum.POWERPC
522 supports_avx512 = torch.cpu._is_avx512_supported()
523 if fp8_kv and not supports_amx and not supports_avx512:
524 raise NotImplementedError(
@@ -525,6 +532,8 @@ def _get_attn_isa(
532 return "neon"
533 elif supports_vxe:
534 return "vxe"
535 + elif supports_vsx:
536 + return "vsx"
537 else:
538 return "vec"
539 else: