Skip to content

GitLab

Unverified Commit ac04dd37 authored by R3hankhan's avatar R3hankhan Committed by GitHub
Browse files

[CPU] Add BF16 Kernel type for s390x (#33788) 


Signed-off-by: default avatarRehan Khan <Rehan.Khan7@ibm.com>
parent 035a6cb0
Hide whitespace changes
Inline Side-by-side
Showing with 9 additions and 0 deletions
csrc/cpu/mla_decode.cpp
View file @ ac04dd37
......@@ -38,6 +38,15 @@ struct KernelVecType<c10::BFloat16> {
using qk_vec_type = vec_op::BF16Vec32;
using v_load_vec_type = vec_op::BF16Vec16;
};
#elif defined(__s390x__)
template <>
struct KernelVecType<c10::BFloat16> {
using qk_load_vec_type = vec_op::BF16Vec16;
using qk_vec_type = vec_op::FP32Vec16;
using v_load_vec_type = vec_op::BF16Vec16;
};
#elif defined(__aarch64__)
template <>
struct KernelVecType<c10::BFloat16> {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment