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 7, 20:15 UTC
[CPU][Perf] Accelerate Attention head for s390x using vector intrinsics (#34434)
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
R3hankhan committed
Feb 24, 2026 at 20:55 UTC
34ce0ffd1f3c7b5e2bcc6073dd1525f266133755
6 files changed
+424
-8
csrc/cpu/cpu_attn.cpp
+4
@@ -16,6 +16,8 @@ torch::Tensor get_scheduler_metadata(
16
isa = cpu_attention::ISA::VEC16;
17
} else if (isa_hint == "neon") {
18
isa = cpu_attention::ISA::NEON;
19
+ } else if (isa_hint == "vxe") {
20
+ isa = cpu_attention::ISA::VXE;
21
} else {
22
TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
23
}
@@ -100,6 +102,8 @@ void cpu_attn_reshape_and_cache(
102
return cpu_attention::ISA::VEC16;
103
} else if (isa == "neon") {
104
return cpu_attention::ISA::NEON;
105
+ } else if (isa == "vxe") {
106
+ return cpu_attention::ISA::VXE;
107
} else {
108
TORCH_CHECK(false, "Invalid ISA type: " + isa);
109
}
csrc/cpu/cpu_attn_impl.hpp
+1
-1
@@ -12,7 +12,7 @@
12
#include "cpu/utils.hpp"
13
14
namespace cpu_attention {
15
-enum class ISA { AMX, VEC, VEC16, NEON };
15
+enum class ISA { AMX, VEC, VEC16, NEON, VXE };
16
17
template <ISA isa, typename scalar_t, int64_t head_dim>
18
class AttentionImpl {};
csrc/cpu/cpu_attn_vxe.hpp
new
+386
@@ -0,0 +1,386 @@
1
+#ifndef CPU_ATTN_VXE_HPP
2
+#define CPU_ATTN_VXE_HPP
3
+
4
+#include "cpu_attn_impl.hpp"
5
+#include <vecintrin.h>
6
+#include <type_traits>
7
+
8
+namespace cpu_attention {
9
+
10
+namespace {
11
+
12
+// s390x Vector = 16 bytes (128 bits)
13
+#define BLOCK_SIZE_ALIGNMENT 32
14
+#define HEAD_SIZE_ALIGNMENT 32
15
+#define MAX_Q_HEAD_NUM_PER_ITER 16
16
+
17
+template <typename kv_cache_t>
18
+FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, __vector float& b0,
19
+ __vector float& b1);
20
+
21
+// [1] Float Specialization
22
+template <>
23
+FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, __vector float& b0,
24
+ __vector float& b1) {
25
+ // Explicitly cast to long long for offset, and float* for pointer
26
+ b0 = vec_xl((long long)0, const_cast<float*>(p));
27
+ b1 = vec_xl((long long)0, const_cast<float*>(p + 4));
28
+}
29
+
30
+// [2] BFloat16 Specialization (Big Endian Fix)
31
+template <>
32
+FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
33
+ __vector float& b0,
34
+ __vector float& b1) {
35
+ // 1. Load 8 BF16s (16 bytes) into one vector
36
+ // Explicit cast to unsigned short* for vec_xl to return vector unsigned short
37
+ __vector unsigned short raw = vec_xl((long long)0, (unsigned short*)p);
38
+
39
+ // 2. Prepare Zero vector
40
+ __vector unsigned short zeros = vec_splat_u16(0);
41
+
42
+ // 3. Merge High/Low to expand BF16 -> Float32
43
+ // On Big Endian, a float is [BF16_bits | 16_zero_bits]
44
+ b0 = (__vector float)vec_mergeh(raw, zeros);
45
+ b1 = (__vector float)vec_mergel(raw, zeros);
46
+}
47
+
48
+template <>
49
+FORCE_INLINE void load_row8_B_as_f32<c10::Half>(const c10::Half* p,
50
+ __vector float& b0,
51
+ __vector float& b1) {
52
+ alignas(16) float tmp[8];
53
+
54
+ // Manual unroll / conversion
55
+ tmp[0] = static_cast<float>(p[0]);
56
+ tmp[1] = static_cast<float>(p[1]);
57
+ tmp[2] = static_cast<float>(p[2]);
58
+ tmp[3] = static_cast<float>(p[3]);
59
+ tmp[4] = static_cast<float>(p[4]);
60
+ tmp[5] = static_cast<float>(p[5]);
61
+ tmp[6] = static_cast<float>(p[6]);
62
+ tmp[7] = static_cast<float>(p[7]);
63
+
64
+ // Explicit arguments for intrinsic: (long long offset, float* ptr)
65
+ b0 = vec_xl((long long)0, (float*)tmp);
66
+ b1 = vec_xl((long long)0, (float*)(tmp + 4));
67
+}
68
+
69
+template <int32_t M, typename kv_cache_t>
70
+FORCE_INLINE void gemm_micro_s390x_Mx8_Ku4(
71
+ const float* __restrict A, // [M x K]
72
+ const kv_cache_t* __restrict B, // [K x 8]
73
+ float* __restrict C, // [M x 8]
74
+ int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
75
+ static_assert(1 <= M && M <= 8, "M must be in [1,8]");
76
+
77
+// Helper macros to unroll codegen for M rows
78
+#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
79
+#define IF_M(i) if constexpr (M > (i))
80
+
81
+ // 1. Define A pointers
82
+#define DECL_A(i) const float* a##i = A + (i) * lda;
83
+ ROWS_APPLY(DECL_A)
84
+#undef DECL_A
85
+
86
+ // 2. Define Accumulators (2 vectors covers 8 columns)
87
+#define DECL_ACC(i) __vector float acc##i##_0, acc##i##_1;
88
+ ROWS_APPLY(DECL_ACC)
89
+#undef DECL_ACC
90
+
91
+ // 3. Initialize Accumulators (Load C or Zero)
92
+#define INIT_ACC(i) \
93
+ IF_M(i) { \
94
+ if (accumulate) { \
95
+ acc##i##_0 = \
96
+ vec_xl((long long)0, const_cast<float*>(C + (i) * ldc + 0)); \
97
+ acc##i##_1 = \
98
+ vec_xl((long long)0, const_cast<float*>(C + (i) * ldc + 4)); \
99
+ } else { \
100
+ acc##i##_0 = vec_splats(0.0f); \
101
+ acc##i##_1 = vec_splats(0.0f); \
102
+ } \
103
+ }
104
+ ROWS_APPLY(INIT_ACC)
105
+#undef INIT_ACC
106
+
107
+ int32_t k = 0;
108
+
109
+ for (; k + 3 < K; k += 4) {
110
+ // Load 4 values of A for each Row M: A[k...k+3]
111
+#define LOAD_A4(i) \
112
+ __vector float a##i##v; \
113
+ IF_M(i) a##i##v = vec_xl((long long)0, const_cast<float*>(a##i + k));
114
+ ROWS_APPLY(LOAD_A4)
115
+#undef LOAD_A4
116
+
117
+ // Helper: FMA for specific lane L of A
118
+ // s390x: vec_madd(b, vec_splat(a, lane), acc)
119
+#define FMAS_LANE(i, aiv, L) \
120
+ IF_M(i) { \
121
+ __vector float a_broad = vec_splat(aiv, L); \
122
+ acc##i##_0 = vec_madd(b0, a_broad, acc##i##_0); \
123
+ acc##i##_1 = vec_madd(b1, a_broad, acc##i##_1); \
124
+ }
125
+
126
+ // Unroll K=0..3
127
+ {
128
+ __vector float b0, b1;
129
+ load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1);
130
+#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
131
+ ROWS_APPLY(STEP_K0)
132
+#undef STEP_K0
133
+ }
134
+ {
135
+ __vector float b0, b1;
136
+ load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1);
137
+#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
138
+ ROWS_APPLY(STEP_K1)
139
+#undef STEP_K1
140
+ }
141
+ {
142
+ __vector float b0, b1;
143
+ load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1);
144
+#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
145
+ ROWS_APPLY(STEP_K2)
146
+#undef STEP_K2
147
+ }
148
+
149
+ {
150
+ __vector float b0, b1;
151
+ load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1);
152
+#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
153
+ ROWS_APPLY(STEP_K3)
154
+#undef STEP_K3
155
+ }
156
+#undef FMAS_LANE
157
+ }
158
+
159
+ for (; k < K; ++k) {
160
+ __vector float b0, b1;
161
+ load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
162
+#define TAIL_ROW(i) \
163
+ IF_M(i) { \
164
+ __vector float ai = vec_splats(*(a##i + k)); \
165
+ acc##i##_0 = vec_madd(b0, ai, acc##i##_0); \
166
+ acc##i##_1 = vec_madd(b1, ai, acc##i##_1); \
167
+ }
168
+ ROWS_APPLY(TAIL_ROW)
169
+#undef TAIL_ROW
170
+ }
171
+
172
+#define STORE_ROW(i) \
173
+ IF_M(i) { \
174
+ vec_xst(acc##i##_0, 0, C + (i) * ldc + 0); \
175
+ vec_xst(acc##i##_1, 0, C + (i) * ldc + 4); \
176
+ }
177
+ ROWS_APPLY(STORE_ROW)
178
+#undef STORE_ROW
179
+
180
+#undef ROWS_APPLY
181
+#undef IF_M
182
+}
183
+
184
+template <int32_t N, typename kv_cache_t>
185
+FORCE_INLINE void gemm_macro_s390x_Mx8_Ku4(const float* __restrict A,
186
+ const kv_cache_t* __restrict B,
187
+ float* __restrict C, int32_t M,
188
+ int32_t K, int64_t lda, int64_t ldb,
189
+ int64_t ldc, bool accumulate) {
190
+ static_assert(N % 8 == 0, "N must be a multiple of 8");
191
+ for (int32_t m = 0; m < M;) {
192
+ int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
193
+ const float* Ab = A + m * lda;
194
+ float* Cb = C + m * ldc;
195
+
196
+ for (int32_t n = 0; n < N; n += 8) {
197
+ const kv_cache_t* Bn = B + n;
198
+ float* Cn = Cb + n;
199
+ switch (mb) {
200
+ case 8:
201
+ gemm_micro_s390x_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
202
+ accumulate);
203
+ break;
204
+ case 4:
205
+ gemm_micro_s390x_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
206
+ accumulate);
207
+ break;
208
+ case 2:
209
+ gemm_micro_s390x_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
210
+ accumulate);
211
+ break;
212
+ default:
213
+ gemm_micro_s390x_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
214
+ accumulate);
215
+ break;
216
+ }
217
+ }
218
+ m += mb;
219
+ }
220
+}
221
+
222
+template <typename kv_cache_t>
223
+class TileGemmS390X {
224
+ public:
225
+ template <AttentionGemmPhase phase, int32_t k_size>
226
+ FORCE_INLINE static void gemm(const int32_t m_size,
227
+ float* __restrict__ a_tile,
228
+ kv_cache_t* __restrict__ b_tile,
229
+ float* __restrict__ c_tile, const int64_t lda,
230
+ const int64_t ldb, const int64_t ldc,
231
+ const int32_t block_size,
232
+ const int32_t dynamic_k_size,
233
+ const bool accum_c) {
234
+ if constexpr (phase == AttentionGemmPhase::QK) {
235
+ gemm_macro_s390x_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
236
+ a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
237
+ } else {
238
+ gemm_macro_s390x_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
239
+ a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
240
+ accum_c);
241
+ }
242
+ }
243
+};
244
+
245
+} // namespace
246
+
247
+template <typename scalar_t, int64_t head_dim>
248
+class AttentionImpl<ISA::VXE, scalar_t, head_dim> {
249
+ public:
250
+ using query_t = scalar_t;
251
+ using q_buffer_t = float;
252
+ using kv_cache_t = scalar_t;
253
+ using logits_buffer_t = float;
254
+ using partial_output_buffer_t = float;
255
+ using prob_buffer_t = float;
256
+
257
+ constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT;
258
+ constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT;
259
+ constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
260
+ constexpr static int64_t HeadDim = head_dim;
261
+ constexpr static ISA ISAType = ISA::VXE;
262
+ constexpr static bool scale_on_logits =
263
+ false; // Scale is applied to Q during copy
264
+
265
+ public:
266
+ AttentionImpl() {}
267
+
268
+ template <template <typename tile_gemm_t> typename attention>
269
+ FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
270
+ attention<TileGemmS390X<kv_cache_t>> attention_iteration;
271
+ attention_iteration(CPU_ATTENTION_PARAMS);
272
+ }
273
+
274
+ // Strides for Memory Layout
275
+ constexpr static int64_t k_cache_token_group_stride(
276
+ const int32_t block_size) {
277
+ return BlockSizeAlignment; // [head_dim, block_size] layout
278
+ }
279
+
280
+ constexpr static int64_t v_cache_token_group_stride(
281
+ const int32_t block_size) {
282
+ return head_dim * BlockSizeAlignment;
283
+ }
284
+
285
+ constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
286
+ return HeadDimAlignment;
287
+ }
288
+
289
+ static void copy_q_heads_tile(scalar_t* __restrict__ src,
290
+ float* __restrict__ q_buffer,
291
+ const int32_t q_num,
292
+ const int32_t q_heads_per_kv,
293
+ const int64_t q_num_stride,
294
+ const int64_t q_head_stride, float scale) {
295
+ __vector float scale_vec = vec_splats(scale);
296
+ constexpr bool is_bf16 = std::is_same<scalar_t, c10::BFloat16>::value;
297
+
298
+ // Process 8 elements at a time (32 bytes of float output)
299
+ for (int32_t i = 0; i < q_num; ++i) {
300
+ for (int32_t h = 0; h < q_heads_per_kv; ++h) {
301
+ scalar_t* curr_src = src + i * q_num_stride + h * q_head_stride;
302
+ float* curr_dst =
303
+ q_buffer + i * q_heads_per_kv * head_dim + h * head_dim;
304
+
305
+ int32_t d = 0;
306
+ for (; d <= head_dim - 8; d += 8) {
307
+ if constexpr (is_bf16) {
308
+ __vector float v0, v1;
309
+ // Reuse our Big-Endian-Safe loader
310
+ load_row8_B_as_f32<scalar_t>(curr_src + d, v0, v1);
311
+
312
+ v0 = vec_mul(v0, scale_vec);
313
+ v1 = vec_mul(v1, scale_vec);
314
+
315
+ vec_xst(v0, 0, curr_dst + d);
316
+ vec_xst(v1, 0, curr_dst + d + 4);
317
+ } else {
318
+ __vector float v0 = vec_xl((long long)0, (float*)curr_src + d);
319
+ __vector float v1 = vec_xl((long long)0, (float*)curr_src + d + 4);
320
+
321
+ v0 = vec_mul(v0, scale_vec);
322
+ v1 = vec_mul(v1, scale_vec);
323
+
324
+ vec_xst(v0, 0, curr_dst + d);
325
+ vec_xst(v1, 0, curr_dst + d + 4);
326
+ }
327
+ }
328
+
329
+ for (; d < head_dim; ++d) {
330
+ float val = static_cast<float>(curr_src[d]);
331
+ curr_dst[d] = val * scale;
332
+ }
333
+ }
334
+ }
335
+ }
336
+
337
+ static void reshape_and_cache(
338
+ const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
339
+ scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
340
+ const int64_t* __restrict__ slot_mapping, const int64_t token_num,
341
+ const int64_t key_token_num_stride, const int64_t value_token_num_stride,
342
+ const int64_t head_num, const int64_t key_head_num_stride,
343
+ const int64_t value_head_num_stride, const int64_t num_blocks,
344
+ const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
345
+ const int64_t block_size, const int64_t block_size_stride) {
346
+#pragma omp parallel for collapse(2)
347
+ for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
348
+ for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
349
+ const int64_t pos = slot_mapping[token_idx];
350
+ if (pos < 0) continue;
351
+
352
+ const int64_t block_idx = pos / block_size;
353
+ const int64_t block_offset = pos % block_size;
354
+
355
+ {
356
+ const scalar_t* key_src = key + token_idx * key_token_num_stride +
357
+ head_idx * key_head_num_stride;
358
+ scalar_t* key_dst = key_cache + block_idx * num_blocks_stride +
359
+ head_idx * cache_head_num_stride + block_offset;
360
+
361
+ for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
362
+ key_dst[j] = key_src[i];
363
+ }
364
+ }
365
+
366
+ {
367
+ const scalar_t* val_src = value + token_idx * value_token_num_stride +
368
+ head_idx * value_head_num_stride;
369
+ scalar_t* val_dst = value_cache + block_idx * num_blocks_stride +
370
+ head_idx * cache_head_num_stride +
371
+ block_offset * head_dim;
372
+
373
+ std::memcpy(val_dst, val_src, sizeof(scalar_t) * head_dim);
374
+ }
375
+ }
376
+ }
377
+ }
378
+};
379
+
380
+} // namespace cpu_attention
381
+
382
+#undef BLOCK_SIZE_ALIGNMENT
383
+#undef HEAD_SIZE_ALIGNMENT
384
+#undef MAX_Q_HEAD_NUM_PER_ITER
385
+
386
+#endif
\ No newline at end of file
csrc/cpu/generate_cpu_attn_dispatch.py
+26
-2
@@ -19,10 +19,11 @@ ISA_TYPES = {
19
"VEC": 1,
20
"VEC16": 2,
21
"NEON": 3,
22
+ "VXE": 4,
23
}
24
25
# ISAs supported for head_dims divisible by 32
25
-ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16"]
26
+ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE"]
27
28
# ISAs supported for head_dims divisible by 16 only
29
ISA_FOR_16 = ["VEC16"]
@@ -118,6 +119,10 @@ def generate_header_file() -> str:
119
#include "cpu_attn_neon.hpp"
120
#endif
121
122
+#ifdef __s390x__
123
+ #include "cpu_attn_vxe.hpp"
124
+#endif
125
+
126
"""
127
128
header += generate_helper_function()
@@ -163,6 +168,25 @@ def generate_header_file() -> str:
168
} \\
169
}()
170
171
+"""
172
+
173
+ # s390x with VXE
174
+ header += """#elif defined(__s390x__)
175
+#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
176
+ [&] { \\
177
+ int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
178
+ switch (encoded_params) { \\
179
+"""
180
+ header += generate_cases_for_isa_group(["VXE", "VEC", "VEC16"])
181
+ header += """
182
+ default: { \\
183
+ TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
184
+ std::to_string(HEAD_DIM) + " isa=" + \\
185
+ std::to_string(static_cast<int>(ISA_TYPE))); \\
186
+ } \\
187
+ } \\
188
+ }()
189
+
190
"""
191
192
# Fallback: VEC and VEC16 only
@@ -182,7 +206,7 @@ def generate_header_file() -> str:
206
} \\
207
}()
208
185
-#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ */
209
+#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */
210
211
#endif // CPU_ATTN_DISPATCH_GENERATED_H
212
"""
vllm/engine/arg_utils.py
+3
-4
@@ -2017,21 +2017,20 @@ class EngineArgs:
2017
)
2018
2019
# Disable chunked prefill and prefix caching for:
2020
- # POWER (ppc64le)/s390x/RISCV CPUs in V1
2020
+ # POWER (ppc64le)/RISCV CPUs in V1
2021
if current_platform.is_cpu() and current_platform.get_cpu_architecture() in (
2022
CpuArchEnum.POWERPC,
2023
- CpuArchEnum.S390X,
2023
CpuArchEnum.RISCV,
2024
):
2025
logger.info(
2026
"Chunked prefill is not supported for POWER, "
2028
- "S390X and RISC-V CPUs; "
2027
+ "and RISC-V CPUs; "
2028
"disabling it for V1 backend."
2029
)
2030
self.enable_chunked_prefill = False
2031
logger.info(
2032
"Prefix caching is not supported for POWER, "
2034
- "S390X and RISC-V CPUs; "
2033
+ "and RISC-V CPUs; "
2034
"disabling it for V1 backend."
2035
)
2036
self.enable_prefix_caching = False
vllm/v1/attention/backends/cpu_attn.py
+4
-1
@@ -25,7 +25,7 @@ from vllm.v1.kv_cache_interface import AttentionSpec, CrossAttentionSpec
25
26
logger = init_logger(__name__)
27
28
-_CPU_ARCH_PREFER_MIXED_BATCH = (CpuArchEnum.X86, CpuArchEnum.ARM)
28
+_CPU_ARCH_PREFER_MIXED_BATCH = (CpuArchEnum.X86, CpuArchEnum.ARM, CpuArchEnum.S390X)
29
30
31
class CPUAttentionBackend(AttentionBackend):
@@ -488,12 +488,15 @@ def _get_attn_isa(
488
return "vec16"
489
supports_amx = torch._C._cpu._is_amx_tile_supported()
490
supports_arm = current_platform.get_cpu_architecture() == CpuArchEnum.ARM
491
+ supports_vxe = current_platform.get_cpu_architecture() == CpuArchEnum.S390X
492
if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0:
493
return "amx"
494
elif block_size % 32 == 0:
495
if supports_arm:
496
# support ARM NEON FMLA and BFMMLA (bf16) for block size 32
497
return "neon"
498
+ elif supports_vxe:
499
+ return "vxe"
500
else:
501
return "vec"
502
else: