[CPU] Speedup LM Head on Arm CPUs (#55352)

Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>

csrc/cpu/dnnl_helper.cpp

@@ -172,7 +172,8 @@ struct hash<MatMulPrimitiveHandler::ClassMatmulCacheKey> {
   size_t operator()(
       const MatMulPrimitiveHandler::ClassMatmulCacheKey& val) const {
     return hash<dnnl_dim_t>()(val.b_n_size) ^ hash<dnnl_dim_t>()(val.b_k_size) ^
-           hash<int>()(static_cast<int>(val.b_type));
+           hash<int>()(static_cast<int>(val.b_type)) ^
+           hash<int>()(static_cast<int>(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_);
   }

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:

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<scalar_t>();
                                });
 
+#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<int64_t>(new MatMulPrimitiveHandler(args));
 }
 
@@ -527,12 +538,22 @@ void onednn_mm(torch::Tensor& c,        // [M, OC], row-major
   MatMulPrimitiveHandler* ptr =
       reinterpret_cast<MatMulPrimitiveHandler*>(handler_tensor.item<int64_t>());
 
+  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<scalar_t>();
 #endif
@@ -563,8 +584,11 @@ void onednn_mm(torch::Tensor& c,        // [M, OC], row-major
     exec_args.a_ptr = a.data_ptr<scalar_t>();
 
 #endif
-    exec_args.c_ptr = c.data_ptr<scalar_t>();
 
     ptr->execute(exec_args);
   });
+
+  if (convert_output) {
+    c.copy_(dnnl_output);
+  }
 }

tests/kernels/test_onednn.py

@@ -15,6 +15,7 @@ NK_FACTORS = [
     (256, 128),
     (4096, 4096),
     (16384, 4096),
+    (64128, 4096),
     (1023, 491),
     (1001, 15),
 ]