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 14, 05:30 UTC
[cpu][fix] Fix onednn_mm crash on consecutive matmuls with same M,K,N and different dtype (#27472)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Fadi Arafeh committed
Oct 24, 2025 at 16:57 UTC
2080b05099b4bf74940a43319babf85c0e0dde57
2 files changed
+9
-4
csrc/cpu/dnnl_helper.cpp
+8
-4
@@ -187,7 +187,8 @@ template <>
187
struct hash<MatMulPrimitiveHandler::ClassMatmulCacheKey> {
188
size_t operator()(
189
const MatMulPrimitiveHandler::ClassMatmulCacheKey& val) const {
190
- return hash<dnnl_dim_t>()(val.b_n_size) ^ hash<dnnl_dim_t>()(val.b_k_size);
190
+ return hash<dnnl_dim_t>()(val.b_n_size) ^ hash<dnnl_dim_t>()(val.b_k_size) ^
191
+ hash<int>()(static_cast<int>(val.b_type));
192
}
193
};
194
@@ -216,7 +217,8 @@ bool operator==(const W8A8MatMulPrimitiveHandler::MSizeCacheKey& l,
217
218
bool operator==(const MatMulPrimitiveHandler::ClassMatmulCacheKey& l,
219
const MatMulPrimitiveHandler::ClassMatmulCacheKey& r) {
219
- return l.b_n_size == r.b_n_size && l.b_k_size == r.b_k_size;
220
+ return l.b_n_size == r.b_n_size && l.b_k_size == r.b_k_size &&
221
+ l.b_type == r.b_type;
222
}
223
224
bool operator==(const MatMulPrimitiveHandler::MSizeCacheKey& l,
@@ -493,8 +495,10 @@ void MatMulPrimitiveHandler::execute(ExecArgs& args) {
495
dnnl::matmul MatMulPrimitiveHandler::get_matmul_cache(
496
const MSizeCacheKey& key) {
497
if (m_size_cache_.get() == nullptr) {
496
- ClassMatmulCacheKey key = {.b_n_size = b_n_size_, .b_k_size = b_k_size_};
497
- m_size_cache_ = get_matul_class_primitive_cache(key, primitive_cache_size_);
498
+ ClassMatmulCacheKey class_key = {
499
+ .b_n_size = b_n_size_, .b_k_size = b_k_size_, .b_type = b_type_};
500
+ m_size_cache_ =
501
+ get_matul_class_primitive_cache(class_key, primitive_cache_size_);
502
}
503
return m_size_cache_->get_or_create(key, [&]() {
504
dnnl::matmul::primitive_desc desc = this->create_primitive_desc(key, false);
csrc/cpu/dnnl_helper.h
+1
@@ -199,6 +199,7 @@ class MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler {
199
struct ClassMatmulCacheKey {
200
dnnl_dim_t b_n_size;
201
dnnl_dim_t b_k_size;
202
+ dnnl::memory::data_type b_type;
203
204
friend bool operator==(const ClassMatmulCacheKey& l,
205
const ClassMatmulCacheKey& r);