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 8, 16:30 UTC
[Fix] [CPU Backend] : Prepack weights for w8a8 oneDNN matmul (#33901)
Signed-off-by: nikhil-arm <nikhil.gupta2@arm.com>
Nikhil Gupta committed
Feb 9, 2026 at 10:04 UTC
caad9f1e01ee04e4f5912d0287031ea3a850f6dc
1 file changed
+10
-2
csrc/cpu/dnnl_helper.cpp
+10
-2
@@ -237,12 +237,20 @@ W8A8MatMulPrimitiveHandler::W8A8MatMulPrimitiveHandler(const Args& args)
237
};
238
dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
239
{b_k_stride_, b_n_stride_});
240
+#ifdef __aarch64__
241
+ // dummy M size for prepacking weights
242
+ // Prepacking weights improves performance and avoid runtime reorders
243
+ constexpr dnnl_dim_t kProbeM = 128;
244
+#else
245
+ constexpr dnnl_dim_t kProbeM = DNNL_RUNTIME_DIM_VAL;
246
+#endif
247
+
248
prepack_weight(args.b_ptr, original_b_md,
249
create_primitive_desc(
242
- MSizeCacheKey{.a_m_size = DNNL_RUNTIME_DIM_VAL,
250
+ MSizeCacheKey{.a_m_size = kProbeM,
251
.use_bias = false,
252
.bias_type = dnnl::memory::data_type::undef},
245
- true)
253
+ /*first_time=*/true)
254
.weights_desc());
255
init_runtime_memory_cache(args);
256
}