diff --git a/csrc/cpu/dnnl_helper.cpp b/csrc/cpu/dnnl_helper.cpp index 14c136dcbb..1ac0401f54 100644 --- a/csrc/cpu/dnnl_helper.cpp +++ b/csrc/cpu/dnnl_helper.cpp @@ -172,7 +172,8 @@ struct hash { size_t operator()( const MatMulPrimitiveHandler::ClassMatmulCacheKey& val) const { return hash()(val.b_n_size) ^ hash()(val.b_k_size) ^ - hash()(static_cast(val.b_type)); + hash()(static_cast(val.b_type)) ^ + hash()(static_cast(val.c_type)); } }; @@ -202,7 +203,7 @@ bool operator==(const W8A8MatMulPrimitiveHandler::MSizeCacheKey& l, bool operator==(const MatMulPrimitiveHandler::ClassMatmulCacheKey& l, const MatMulPrimitiveHandler::ClassMatmulCacheKey& r) { return l.b_n_size == r.b_n_size && l.b_k_size == r.b_k_size && - l.b_type == r.b_type; + l.b_type == r.b_type && l.c_type == r.c_type; } bool operator==(const MatMulPrimitiveHandler::MSizeCacheKey& l, @@ -482,8 +483,10 @@ void MatMulPrimitiveHandler::execute(ExecArgs& args) { dnnl::matmul MatMulPrimitiveHandler::get_matmul_cache( const MSizeCacheKey& key) { if (m_size_cache_.get() == nullptr) { - ClassMatmulCacheKey class_key = { - .b_n_size = b_n_size_, .b_k_size = b_k_size_, .b_type = b_type_}; + ClassMatmulCacheKey class_key = {.b_n_size = b_n_size_, + .b_k_size = b_k_size_, + .b_type = b_type_, + .c_type = c_type_}; m_size_cache_ = get_matul_class_primitive_cache(class_key, primitive_cache_size_); } diff --git a/csrc/cpu/dnnl_helper.h b/csrc/cpu/dnnl_helper.h index b841ed7305..05fb116ecd 100644 --- a/csrc/cpu/dnnl_helper.h +++ b/csrc/cpu/dnnl_helper.h @@ -176,6 +176,7 @@ class MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler { dnnl_dim_t b_n_size; dnnl_dim_t b_k_size; dnnl::memory::data_type b_type; + dnnl::memory::data_type c_type; friend bool operator==(const ClassMatmulCacheKey& l, const ClassMatmulCacheKey& r); @@ -203,6 +204,8 @@ class MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler { public: MatMulPrimitiveHandler(const Args& args); + dnnl::memory::data_type get_output_type() const { return c_type_; } + void execute(ExecArgs& args); private: diff --git a/csrc/cpu/dnnl_kernels.cpp b/csrc/cpu/dnnl_kernels.cpp index 6dda092961..e81d365628 100644 --- a/csrc/cpu/dnnl_kernels.cpp +++ b/csrc/cpu/dnnl_kernels.cpp @@ -513,6 +513,17 @@ int64_t create_onednn_mm_handler(const torch::Tensor& b, args.ab_type = get_dnnl_type(); }); +#ifdef VLLM_USE_ACL + // TODO(fadara01): this is a hot-fix, remove once problem is addressed in the + // AArch64 stack. hybrid kernels in ACL which are the fastest for LM head + // layers are only enabled for BF16 x BF16 -> FP32 problem desc as they can't + // convert FP32->BF16 internally. In this case do: oneDNN matmul with + // bf16:bf16:fp32 then cast output down to bf16. + if (b.scalar_type() == at::kBFloat16 && b.size(1) >= 64000) { + args.c_type = dnnl::memory::data_type::f32; + } +#endif + return reinterpret_cast(new MatMulPrimitiveHandler(args)); } @@ -527,12 +538,22 @@ void onednn_mm(torch::Tensor& c, // [M, OC], row-major MatMulPrimitiveHandler* ptr = reinterpret_cast(handler_tensor.item()); + torch::Tensor dnnl_output = c; + // AArch64 specific case where we do a bf16 x bf16 -> fp32 matmul + const bool convert_output = + ptr->get_output_type() == dnnl::memory::data_type::f32 && + c.scalar_type() != at::kFloat; + if (convert_output) { + dnnl_output = torch::empty(c.sizes(), c.options().dtype(at::kFloat)); + } + // ACL matmuls expect contiguous source tensors #ifdef VLLM_USE_ACL torch::Tensor a_contig = a.contiguous(); #endif MatMulPrimitiveHandler::ExecArgs exec_args; + exec_args.c_ptr = dnnl_output.data_ptr(); #ifdef VLLM_USE_ACL exec_args.a_m_size = a_contig.size(0); @@ -548,7 +569,7 @@ void onednn_mm(torch::Tensor& c, // [M, OC], row-major #ifdef VLLM_USE_ACL // ACL matmuls in oneDNN do not support a bias. // We handle a matmul with bias by doing: c = bias; c += matmul(a, b) - c.copy_(bias.value()); + dnnl_output.copy_(bias.value()); #else exec_args.bias_ptr = bias->data_ptr(); #endif @@ -563,8 +584,11 @@ void onednn_mm(torch::Tensor& c, // [M, OC], row-major exec_args.a_ptr = a.data_ptr(); #endif - exec_args.c_ptr = c.data_ptr(); ptr->execute(exec_args); }); + + if (convert_output) { + c.copy_(dnnl_output); + } } diff --git a/tests/kernels/test_onednn.py b/tests/kernels/test_onednn.py index 4e681ca6ac..4e41791a54 100644 --- a/tests/kernels/test_onednn.py +++ b/tests/kernels/test_onednn.py @@ -15,6 +15,7 @@ NK_FACTORS = [ (256, 128), (4096, 4096), (16384, 4096), + (64128, 4096), (1023, 491), (1001, 15), ]