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 8, 19:00 UTC
[CPU] Add BF16 Kernel type for s390x (#33788)
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com>
R3hankhan committed
Feb 6, 2026 at 10:27 UTC
ac04dd374f996f8df960933fa076bb5ea53c0c2a
1 file changed
+9
csrc/cpu/mla_decode.cpp
+9
@@ -38,6 +38,15 @@ struct KernelVecType<c10::BFloat16> {
38
using qk_vec_type = vec_op::BF16Vec32;
39
using v_load_vec_type = vec_op::BF16Vec16;
40
};
41
+
42
+#elif defined(__s390x__)
43
+template <>
44
+struct KernelVecType<c10::BFloat16> {
45
+ using qk_load_vec_type = vec_op::BF16Vec16;
46
+ using qk_vec_type = vec_op::FP32Vec16;
47
+ using v_load_vec_type = vec_op::BF16Vec16;
48
+};
49
+
50
#elif defined(__aarch64__)
51
template <>
52
struct KernelVecType<c10::BFloat16> {