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 7, 05:15 UTC
[Perf] Use dummy M for weight prepacking on x86 (#35890)
Signed-off-by: Li, Tianmu <tianmu.li@intel.com>
Tianmu Li committed
Mar 4, 2026 at 20:56 UTC
8e7820131ee8d0295e6a533d745f6ca8085baec9
1 file changed
+12
-17
csrc/cpu/dnnl_helper.cpp
+12
-17
@@ -237,13 +237,10 @@ 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__
240
+
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
244
245
prepack_weight(args.b_ptr, original_b_md,
246
create_primitive_desc(
@@ -411,21 +408,19 @@ MatMulPrimitiveHandler::MatMulPrimitiveHandler(const Args& args)
408
dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
409
{b_k_stride_, b_n_stride_});
410
411
+ // dummy M size for prepacking weights
412
+ // Prepacking weights improves performance and avoid runtime reorders
413
+ constexpr dnnl_dim_t kProbeM = 128;
414
+
415
prepack_weight(args.b_ptr, original_b_md,
416
create_primitive_desc(
416
- MSizeCacheKey{
417
-#ifdef VLLM_USE_ACL
418
- // Arm Compute Library (ACL) backend for oneDNN does
419
- // not support runtime
420
- // dimensions, so we set M to a default value
421
- .a_m_size = 128,
422
- .a_m_stride = b_k_size_,
423
-#else
424
- .a_m_size = DNNL_RUNTIME_DIM_VAL,
425
- .a_m_stride = DNNL_RUNTIME_DIM_VAL,
426
-#endif
427
- .use_bias = false,
428
- .bias_type = dnnl::memory::data_type::undef},
417
+ MSizeCacheKey{// Use a concrete M so oneDNN's kernel
418
+ // selector can choose an optimally blocked
419
+ // weight layout.
420
+ .a_m_size = kProbeM,
421
+ .a_m_stride = b_k_size_,
422
+ .use_bias = false,
423
+ .bias_type = dnnl::memory::data_type::undef},
424
true)
425
.weights_desc());
426
init_runtime_memory_cache(args);