[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>
Fadi Arafeh committed
Sep 11, 2026 at 09:21 UTC
e3f755b732316df4db8437dcb973b9fa28bfd656
4 files changed
+37
-6
csrc/cpu/dnnl_helper.cpp
+7
-4
@@ -172,7 +172,8 @@ struct hash<MatMulPrimitiveHandler::ClassMatmulCacheKey> {
172
size_t operator()(
173
const MatMulPrimitiveHandler::ClassMatmulCacheKey& val) const {
174
return hash<dnnl_dim_t>()(val.b_n_size) ^ hash<dnnl_dim_t>()(val.b_k_size) ^
175
- hash<int>()(static_cast<int>(val.b_type));
175
+ hash<int>()(static_cast<int>(val.b_type)) ^
176
+ hash<int>()(static_cast<int>(val.c_type));
177
}
178
};
179
@@ -202,7 +203,7 @@ bool operator==(const W8A8MatMulPrimitiveHandler::MSizeCacheKey& l,
203
bool operator==(const MatMulPrimitiveHandler::ClassMatmulCacheKey& l,
204
const MatMulPrimitiveHandler::ClassMatmulCacheKey& r) {
205
return l.b_n_size == r.b_n_size && l.b_k_size == r.b_k_size &&
205
- l.b_type == r.b_type;
206
+ l.b_type == r.b_type && l.c_type == r.c_type;
207
}
208
209
bool operator==(const MatMulPrimitiveHandler::MSizeCacheKey& l,
@@ -482,8 +483,10 @@ void MatMulPrimitiveHandler::execute(ExecArgs& args) {
483
dnnl::matmul MatMulPrimitiveHandler::get_matmul_cache(
484
const MSizeCacheKey& key) {
485
if (m_size_cache_.get() == nullptr) {
485
- ClassMatmulCacheKey class_key = {
486
- .b_n_size = b_n_size_, .b_k_size = b_k_size_, .b_type = b_type_};
486
+ ClassMatmulCacheKey class_key = {.b_n_size = b_n_size_,
487
+ .b_k_size = b_k_size_,
488
+ .b_type = b_type_,
489
+ .c_type = c_type_};
490
m_size_cache_ =
491
get_matul_class_primitive_cache(class_key, primitive_cache_size_);
492
}
csrc/cpu/dnnl_helper.h
+3
@@ -176,6 +176,7 @@ class MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler {
176
dnnl_dim_t b_n_size;
177
dnnl_dim_t b_k_size;
178
dnnl::memory::data_type b_type;
179
+ dnnl::memory::data_type c_type;
180
181
friend bool operator==(const ClassMatmulCacheKey& l,
182
const ClassMatmulCacheKey& r);
@@ -203,6 +204,8 @@ class MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler {
204
public:
205
MatMulPrimitiveHandler(const Args& args);
206
207
+ dnnl::memory::data_type get_output_type() const { return c_type_; }
208
+
209
void execute(ExecArgs& args);
210
211
private:
csrc/cpu/dnnl_kernels.cpp
+26
-2
@@ -513,6 +513,17 @@ int64_t create_onednn_mm_handler(const torch::Tensor& b,
513
args.ab_type = get_dnnl_type<scalar_t>();
514
});
515
516
+#ifdef VLLM_USE_ACL
517
+ // TODO(fadara01): this is a hot-fix, remove once problem is addressed in the
518
+ // AArch64 stack. hybrid kernels in ACL which are the fastest for LM head
519
+ // layers are only enabled for BF16 x BF16 -> FP32 problem desc as they can't
520
+ // convert FP32->BF16 internally. In this case do: oneDNN matmul with
521
+ // bf16:bf16:fp32 then cast output down to bf16.
522
+ if (b.scalar_type() == at::kBFloat16 && b.size(1) >= 64000) {
523
+ args.c_type = dnnl::memory::data_type::f32;
524
+ }
525
+#endif
526
+
527
return reinterpret_cast<int64_t>(new MatMulPrimitiveHandler(args));
528
}
529
@@ -527,12 +538,22 @@ void onednn_mm(torch::Tensor& c, // [M, OC], row-major
538
MatMulPrimitiveHandler* ptr =
539
reinterpret_cast<MatMulPrimitiveHandler*>(handler_tensor.item<int64_t>());
540
541
+ torch::Tensor dnnl_output = c;
542
+ // AArch64 specific case where we do a bf16 x bf16 -> fp32 matmul
543
+ const bool convert_output =
544
+ ptr->get_output_type() == dnnl::memory::data_type::f32 &&
545
+ c.scalar_type() != at::kFloat;
546
+ if (convert_output) {
547
+ dnnl_output = torch::empty(c.sizes(), c.options().dtype(at::kFloat));
548
+ }
549
+
550
// ACL matmuls expect contiguous source tensors
551
#ifdef VLLM_USE_ACL
552
torch::Tensor a_contig = a.contiguous();
553
#endif
554
555
MatMulPrimitiveHandler::ExecArgs exec_args;
556
+ exec_args.c_ptr = dnnl_output.data_ptr();
557
558
#ifdef VLLM_USE_ACL
559
exec_args.a_m_size = a_contig.size(0);
@@ -548,7 +569,7 @@ void onednn_mm(torch::Tensor& c, // [M, OC], row-major
569
#ifdef VLLM_USE_ACL
570
// ACL matmuls in oneDNN do not support a bias.
571
// We handle a matmul with bias by doing: c = bias; c += matmul(a, b)
551
- c.copy_(bias.value());
572
+ dnnl_output.copy_(bias.value());
573
#else
574
exec_args.bias_ptr = bias->data_ptr<scalar_t>();
575
#endif
@@ -563,8 +584,11 @@ void onednn_mm(torch::Tensor& c, // [M, OC], row-major
584
exec_args.a_ptr = a.data_ptr<scalar_t>();
585
586
#endif
566
- exec_args.c_ptr = c.data_ptr<scalar_t>();
587
588
ptr->execute(exec_args);
589
});
590
+
591
+ if (convert_output) {
592
+ c.copy_(dnnl_output);
593
+ }
594
}
tests/kernels/test_onednn.py
+1
@@ -15,6 +15,7 @@ NK_FACTORS = [
15
(256, 128),
16
(4096, 4096),
17
(16384, 4096),
18
+ (64128, 4096),
19
(1023, 491),
20
(1001, 15),
21
]